1use crate::{Edge, EdgeKind, FactSet, Node, NodeKind, Provenance, Span};
16
17pub(crate) const EXTRACT_VERSION: u32 = 9
40 + if cfg!(feature = "pdf-text") { 100 } else { 0 }
41 + if cfg!(feature = "image-ocr") { 200 } else { 0 }
42 + if cfg!(feature = "image-vision") {
43 400
44 } else {
45 0
46 }
47 + if cfg!(feature = "audio-transcribe") {
48 800
49 } else {
50 0
51 };
52
53const MAX_CONTENT: usize = 1500;
57
58#[cfg(feature = "pdf-text")]
61const MAX_PDF_BYTES: usize = 20 * 1024 * 1024;
62
63#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
65const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024;
66
67#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
71const MAX_IMAGE_PIXELS: u64 = 4096 * 4096;
72
73#[cfg(feature = "image-vision")]
77const MIN_OCR_WORDS: usize = 8;
78
79#[cfg(feature = "audio-transcribe")]
82const MAX_AUDIO_BYTES: usize = 50 * 1024 * 1024;
83
84pub trait Extractor {
86 fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet;
91
92 fn env_tag(&self) -> u64 {
99 media_env_tag()
100 }
101}
102
103#[allow(clippy::struct_excessive_bools)]
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub struct IngestConfig {
113 pub prose: bool,
115 pub pdf: bool,
117 pub ocr: bool,
119 pub vision: bool,
121 pub audio: bool,
123}
124
125impl Default for IngestConfig {
126 fn default() -> Self {
127 Self {
128 prose: true,
129 pdf: true,
130 ocr: true,
131 vision: true,
132 audio: true,
133 }
134 }
135}
136
137impl IngestConfig {
138 fn disabled_bits(self) -> u64 {
143 u64::from(!self.prose)
144 | (u64::from(!self.pdf) << 1)
145 | (u64::from(!self.ocr) << 2)
146 | (u64::from(!self.vision) << 3)
147 | (u64::from(!self.audio) << 4)
148 }
149}
150
151#[derive(Debug, Clone, Copy, Default)]
157pub struct Registry {
158 pub ingest: IngestConfig,
160}
161
162impl Registry {
163 #[must_use]
165 pub fn new(ingest: IngestConfig) -> Self {
166 Self { ingest }
167 }
168}
169
170impl Extractor for Registry {
171 fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
172 let mut facts = extract_facts(path, blob_id, bytes, self.ingest);
173 crate::markers::augment(&mut facts, path, blob_id, bytes);
174 facts
175 }
176
177 fn env_tag(&self) -> u64 {
178 let media = media_env_tag();
179 let disabled = self.ingest.disabled_bits();
180 if disabled == 0 {
181 media
183 } else {
184 let mut h = 0xcbf2_9ce4_8422_2325u64;
189 for b in media
190 .to_le_bytes()
191 .into_iter()
192 .chain(disabled.to_le_bytes())
193 {
194 h ^= u64::from(b);
195 h = h.wrapping_mul(0x0000_0100_0000_01b3);
196 }
197 h
198 }
199 }
200}
201
202fn extract_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
206 if crate::config_keys::is_config_path(path) {
209 return config_facts(path, blob_id, bytes, ingest);
210 }
211 if is_dockerfile(path) {
214 return dockerfile_facts(path, blob_id, bytes, ingest);
215 }
216 let ext = extension(path);
217 match ext.as_deref() {
218 Some("rs") => rust_facts(path, blob_id, bytes, ingest),
220 Some(ext) => tag_facts(path, blob_id, bytes, ext, ingest).unwrap_or_else(|| {
224 FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest))
225 }),
226 None => FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest)),
227 }
228}
229
230fn extension(path: &str) -> Option<String> {
233 let name = path.rsplit('/').next().unwrap_or(path);
234 name.rsplit_once('.')
235 .map(|(_, ext)| ext.to_ascii_lowercase())
236}
237
238fn file_key(path: &str) -> String {
240 format!("file:{path}")
241}
242
243fn file_node(
247 path: &str,
248 blob_id: &str,
249 bytes: &[u8],
250 lang: Option<&str>,
251 ingest: IngestConfig,
252) -> Node {
253 let name = path.rsplit('/').next().unwrap_or(path).to_owned();
254 let lines = bytes
255 .iter()
256 .fold(0usize, |n, &b| n + usize::from(b == b'\n'));
257 let end = u32::try_from(bytes.len()).unwrap_or(u32::MAX);
258 let mut meta = serde_json::json!({ "bytes": bytes.len(), "lines": lines });
259 let content = if ingest.prose && is_prose(path) {
264 cap_content(&String::from_utf8_lossy(bytes))
265 } else if let Some(text) = ingest.pdf.then(|| pdf_content(path, bytes)).flatten() {
266 cap_content(&text)
267 } else if let Some(text) = image_content(path, bytes, ingest) {
268 cap_content(&text)
269 } else if let Some(text) = audio_content(path, bytes, ingest) {
270 cap_content(&text)
271 } else {
272 String::new()
273 };
274 if !content.is_empty() {
275 meta["content"] = serde_json::Value::from(content);
276 }
277 Node {
278 key: file_key(path),
279 kind: NodeKind::File,
280 name,
281 path: Some(path.to_owned()),
282 lang: lang.map(ToOwned::to_owned),
283 blob_hash: Some(blob_id.to_owned()),
284 span: Some(Span::new(0, end)),
285 provenance: Provenance::Derived,
286 meta,
287 }
288}
289
290fn config_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
296 let mut facts = FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
297 let file = file_key(path);
298 let mut by_key: std::collections::BTreeMap<String, String> = std::collections::BTreeMap::new();
301 for ck in crate::config_keys::flatten(path, bytes) {
302 by_key.insert(ck.key, ck.value);
303 }
304 for (key, value) in by_key {
305 let node_key = format!("cfgkey:{path}#{key}");
306 let value = if crate::config_keys::is_secret_key(&key) {
309 "<redacted>".to_owned()
310 } else {
311 value
312 };
313 let mut node = Node::new(
314 node_key.clone(),
315 NodeKind::Other(crate::config_keys::KIND.into()),
316 key.clone(),
317 );
318 node.path = Some(path.to_owned());
319 node.blob_hash = Some(blob_id.to_owned());
320 node.meta = serde_json::json!({ "key": key, "value": value });
321 facts = facts.with_node(node).with_edge(Edge::derived(
322 file.clone(),
323 node_key,
324 EdgeKind::Contains,
325 ));
326 }
327 facts
328}
329
330pub(crate) const IMAGE_REF_KIND: &str = "image_ref";
334
335fn is_dockerfile(path: &str) -> bool {
338 let base = path.rsplit('/').next().unwrap_or(path).to_ascii_lowercase();
339 base == "dockerfile"
340 || base == "containerfile"
341 || base.starts_with("dockerfile.")
342 || base.ends_with(".dockerfile")
343}
344
345fn dockerfile_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
350 let mut facts = FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
351 let file = file_key(path);
352 let text = String::from_utf8_lossy(bytes);
353 let mut stages: std::collections::HashSet<String> = std::collections::HashSet::new();
354 let mut idx = 0usize;
355 for line in text.lines() {
356 let Some(rest) = strip_from_prefix(line.trim()) else {
357 continue;
358 };
359 let (image, stage) = parse_from(rest);
360 let is_internal_stage = stages.contains(&image.to_ascii_lowercase());
364 if let Some(s) = stage {
365 stages.insert(s.to_ascii_lowercase());
366 }
367 if image.is_empty() || image.eq_ignore_ascii_case("scratch") || is_internal_stage {
370 continue;
371 }
372 let (name, tag, digest) = split_image(image);
373 let node_key = format!("imageref:{path}#{idx}");
374 idx += 1;
375 let mut node = Node::new(
376 node_key.clone(),
377 NodeKind::Other(IMAGE_REF_KIND.into()),
378 image.to_owned(),
379 );
380 node.path = Some(path.to_owned());
381 node.blob_hash = Some(blob_id.to_owned());
382 node.meta = serde_json::json!({ "image": name, "tag": tag, "digest": digest });
383 facts = facts.with_node(node).with_edge(Edge::derived(
384 file.clone(),
385 node_key,
386 EdgeKind::References,
387 ));
388 }
389 facts
390}
391
392fn strip_from_prefix(line: &str) -> Option<&str> {
394 let b = line.as_bytes();
395 (b.len() >= 5 && b[..4].eq_ignore_ascii_case(b"from") && b[4].is_ascii_whitespace())
396 .then(|| line[5..].trim_start())
397}
398
399fn parse_from(rest: &str) -> (&str, Option<&str>) {
403 let image = rest
404 .split_whitespace()
405 .find(|t| !t.starts_with("--"))
406 .unwrap_or("");
407 let mut toks = rest.split_whitespace();
408 let mut stage = None;
409 while let Some(t) = toks.next() {
410 if t.eq_ignore_ascii_case("as") {
411 stage = toks.next();
412 break;
413 }
414 }
415 (image, stage)
416}
417
418fn split_image(image: &str) -> (String, Option<String>, Option<String>) {
422 if let Some((name, digest)) = image.split_once('@') {
423 return (name.to_owned(), None, Some(digest.to_owned()));
424 }
425 let seg = image.rfind('/').map_or(0, |i| i + 1);
426 if let Some(colon) = image[seg..].find(':') {
427 let at = seg + colon;
428 return (
429 image[..at].to_owned(),
430 Some(image[at + 1..].to_owned()),
431 None,
432 );
433 }
434 (image.to_owned(), None, None)
435}
436
437fn doc_comment_body(raw: &str) -> Option<String> {
441 let t = raw.trim();
442 if t.starts_with("//!") || (t.starts_with("///") && !t.starts_with("////")) {
443 return Some(t[3..].trim().to_owned());
444 }
445 if (t.starts_with("/**") || t.starts_with("/*!")) && t.ends_with("*/") {
446 let end = t.len() - 2;
450 let inner = if end >= 3 { &t[3..end] } else { "" };
451 let cleaned: Vec<&str> = inner
452 .lines()
453 .map(|l| l.trim().trim_start_matches('*').trim())
454 .filter(|l| !l.is_empty())
455 .collect();
456 return Some(cleaned.join(" "));
457 }
458 None
459}
460
461#[cfg(feature = "pdf-text")]
469fn pdf_content(path: &str, bytes: &[u8]) -> Option<String> {
470 if extension(path).as_deref() != Some("pdf") || bytes.len() > MAX_PDF_BYTES {
471 return None;
472 }
473 let owned = bytes.to_vec();
474 let text = std::panic::catch_unwind(move || pdf_extract::extract_text_from_mem(&owned).ok())
475 .ok()
476 .flatten()?;
477 (!text.trim().is_empty()).then_some(text)
478}
479
480#[cfg(not(feature = "pdf-text"))]
482fn pdf_content(_path: &str, _bytes: &[u8]) -> Option<String> {
483 None
484}
485
486#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
496fn image_content(path: &str, bytes: &[u8], ingest: IngestConfig) -> Option<String> {
497 if !is_image(path) || bytes.len() > MAX_IMAGE_BYTES {
498 return None;
499 }
500 let ocr = if ingest.ocr { ocr_content(bytes) } else { None };
506 let sparse = ocr
507 .as_deref()
508 .is_none_or(|t| t.split_whitespace().count() < min_ocr_words());
509 let vision = if ingest.vision && sparse {
510 vlm_content(bytes)
511 } else {
512 None
513 };
514 match (ocr, vision) {
515 (Some(o), Some(v)) => Some(format!("{o}\n\n{v}")),
516 (Some(o), None) => Some(o),
517 (None, Some(v)) => Some(v),
518 (None, None) => None,
519 }
520}
521
522#[cfg(not(any(feature = "image-ocr", feature = "image-vision")))]
524fn image_content(_path: &str, _bytes: &[u8], _ingest: IngestConfig) -> Option<String> {
525 None
526}
527
528#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
531fn min_ocr_words() -> usize {
532 #[cfg(feature = "image-vision")]
533 {
534 MIN_OCR_WORDS
535 }
536 #[cfg(not(feature = "image-vision"))]
537 {
538 usize::MAX
539 }
540}
541
542#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
544fn is_image(path: &str) -> bool {
545 matches!(extension(path).as_deref(), Some("png" | "jpg" | "jpeg"))
546}
547
548#[cfg(feature = "audio-transcribe")]
557fn audio_content(path: &str, bytes: &[u8], ingest: IngestConfig) -> Option<String> {
558 if !ingest.audio || !is_audio(path) || bytes.len() > MAX_AUDIO_BYTES {
559 return None;
560 }
561 asr_content(bytes)
562}
563
564#[cfg(not(feature = "audio-transcribe"))]
567fn audio_content(_path: &str, _bytes: &[u8], _ingest: IngestConfig) -> Option<String> {
568 None
569}
570
571#[cfg(feature = "audio-transcribe")]
574fn is_audio(path: &str) -> bool {
575 matches!(extension(path).as_deref(), Some("wav" | "mp3" | "flac"))
576}
577
578#[cfg(feature = "audio-transcribe")]
583fn asr_content(bytes: &[u8]) -> Option<String> {
584 use rto_llama::Engine as _;
585
586 let engine = asr_engine()?;
587 let completion = engine
588 .chat(&rto_llama::ChatRequest {
589 model: ASR_MODEL.to_owned(),
590 messages: vec![rto_llama::Message {
591 role: "user".to_owned(),
592 content: "Transcribe this audio recording. Output only the spoken words, verbatim."
593 .to_owned(),
594 }],
595 images: Vec::new(),
596 audio: vec![bytes.to_vec()],
597 temperature: 0.0,
598 max_tokens: 512,
599 })
600 .ok()?;
601 let text = completion.content.trim();
602 (!text.is_empty()).then(|| text.to_owned())
603}
604
605#[cfg(feature = "audio-transcribe")]
607const ASR_MODEL: &str = "voxtral-mini-3b";
608
609#[cfg(feature = "audio-transcribe")]
613fn asr_engine() -> Option<&'static rto_llama::llama::LlamaEngine> {
614 use std::sync::OnceLock;
615 static ENGINE: OnceLock<Option<rto_llama::llama::LlamaEngine>> = OnceLock::new();
616 ENGINE
617 .get_or_init(|| {
618 let dir = crate::models::model_dir(ASR_MODEL);
619 let (gguf, mmproj) = (dir.join("model.gguf"), dir.join("mmproj.gguf"));
620 if !gguf.exists() || !mmproj.exists() {
621 return None;
622 }
623 rto_llama::llama::LlamaEngine::new(
624 vec![rto_llama::llama::Served {
625 name: ASR_MODEL.to_owned(),
626 path: gguf,
627 mmproj: Some(mmproj),
628 }],
629 0,
630 )
631 .ok()
632 })
633 .as_ref()
634}
635
636#[cfg(any(feature = "image-ocr", feature = "image-vision"))]
641fn image_dimensions_ok(bytes: &[u8]) -> bool {
642 let Ok(reader) = image::ImageReader::new(std::io::Cursor::new(bytes)).with_guessed_format()
643 else {
644 return false;
645 };
646 match reader.into_dimensions() {
647 Ok((w, h)) => u64::from(w) * u64::from(h) <= MAX_IMAGE_PIXELS,
648 Err(_) => false,
649 }
650}
651
652#[cfg(feature = "image-ocr")]
656fn ocr_content(bytes: &[u8]) -> Option<String> {
657 let dir = crate::models::model_dir("ocrs-text");
658 let detection = dir.join("text-detection.rten");
659 let recognition = dir.join("text-recognition.rten");
660 if !detection.exists() || !recognition.exists() || !image_dimensions_ok(bytes) {
661 return None;
663 }
664 let text = std::panic::catch_unwind(|| run_ocr(&detection, &recognition, bytes))
667 .ok()
668 .flatten()?;
669 (!text.trim().is_empty()).then_some(text)
670}
671
672#[cfg(all(feature = "image-vision", not(feature = "image-ocr")))]
677fn ocr_content(_bytes: &[u8]) -> Option<String> {
678 None
679}
680
681#[cfg(feature = "image-ocr")]
684fn run_ocr(
685 detection: &std::path::Path,
686 recognition: &std::path::Path,
687 bytes: &[u8],
688) -> Option<String> {
689 use ocrs::{ImageSource, OcrEngine, OcrEngineParams};
690
691 let detection_model = rten::Model::load_file(detection).ok()?;
692 let recognition_model = rten::Model::load_file(recognition).ok()?;
693 let engine = OcrEngine::new(OcrEngineParams {
694 detection_model: Some(detection_model),
695 recognition_model: Some(recognition_model),
696 ..Default::default()
697 })
698 .ok()?;
699
700 let img = image::load_from_memory(bytes).ok()?.into_rgb8();
701 let source = ImageSource::from_bytes(img.as_raw(), img.dimensions()).ok()?;
702 let input = engine.prepare_input(source).ok()?;
703 engine.get_text(&input).ok()
704}
705
706#[cfg(feature = "image-vision")]
713fn vlm_content(bytes: &[u8]) -> Option<String> {
714 use rto_llama::Engine as _;
715
716 if !image_dimensions_ok(bytes) {
717 return None;
718 }
719 let engine = vlm_engine()?;
720 let completion = engine
721 .chat(&rto_llama::ChatRequest {
722 model: VLM_MODEL.to_owned(),
723 messages: vec![rto_llama::Message {
724 role: "user".to_owned(),
725 content: "Describe this image in one or two sentences.".to_owned(),
726 }],
727 images: vec![bytes.to_vec()],
728 audio: Vec::new(),
729 temperature: 0.0,
730 max_tokens: 128,
731 })
732 .ok()?;
733 let text = completion.content.trim();
734 (!text.is_empty()).then(|| text.to_owned())
735}
736
737#[cfg(feature = "image-vision")]
739const VLM_MODEL: &str = "smolvlm-500m-gguf";
740
741#[cfg(feature = "image-vision")]
745fn vlm_engine() -> Option<&'static rto_llama::llama::LlamaEngine> {
746 use std::sync::OnceLock;
747 static ENGINE: OnceLock<Option<rto_llama::llama::LlamaEngine>> = OnceLock::new();
748 ENGINE
749 .get_or_init(|| {
750 let dir = crate::models::model_dir(VLM_MODEL);
751 let (gguf, mmproj) = (dir.join("model.gguf"), dir.join("mmproj.gguf"));
752 if !gguf.exists() || !mmproj.exists() {
753 return None;
754 }
755 rto_llama::llama::LlamaEngine::new(
756 vec![rto_llama::llama::Served {
757 name: VLM_MODEL.to_owned(),
758 path: gguf,
759 mmproj: Some(mmproj),
760 }],
761 0,
762 )
763 .ok()
764 })
765 .as_ref()
766}
767
768#[cfg(all(feature = "image-ocr", not(feature = "image-vision")))]
771fn vlm_content(_bytes: &[u8]) -> Option<String> {
772 None
773}
774
775#[cfg(any(
785 feature = "image-ocr",
786 feature = "image-vision",
787 feature = "audio-transcribe"
788))]
789pub(crate) fn media_env_tag() -> u64 {
790 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
791 let mut any = false;
792 #[cfg(feature = "image-ocr")]
793 {
794 any |= fold_installed_model(&mut hash, "ocrs-text");
795 }
796 #[cfg(feature = "image-vision")]
797 {
798 any |= fold_installed_model(&mut hash, "smolvlm-500m-gguf");
799 }
800 #[cfg(feature = "audio-transcribe")]
801 {
802 any |= fold_installed_model(&mut hash, "voxtral-mini-3b");
803 }
804 if any { hash | 1 } else { 0 }
805}
806
807#[cfg(any(
811 feature = "image-ocr",
812 feature = "image-vision",
813 feature = "audio-transcribe"
814))]
815fn fold_installed_model(hash: &mut u64, name: &str) -> bool {
816 let Some(variant) = crate::models::find(name)
817 .and_then(|spec| spec.variant_for(crate::models::Platform::host()))
818 else {
819 return false;
820 };
821 let dir = crate::models::model_dir(name);
822 if !variant.files.iter().all(|f| dir.join(f.name).exists()) {
823 return false;
824 }
825 for file in variant.files {
826 for b in file.sha256.bytes() {
827 *hash ^= u64::from(b);
828 *hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
829 }
830 }
831 true
832}
833
834#[cfg(not(any(
836 feature = "image-ocr",
837 feature = "image-vision",
838 feature = "audio-transcribe"
839)))]
840pub(crate) fn media_env_tag() -> u64 {
841 0
842}
843
844fn is_prose(path: &str) -> bool {
846 matches!(
847 extension(path).as_deref(),
848 Some("md" | "markdown" | "txt" | "rst" | "adoc")
849 )
850}
851
852fn cap_content(text: &str) -> String {
855 let mut out = String::with_capacity(text.len().min(MAX_CONTENT));
856 let mut chars = 0usize;
859 let mut last_was_space = true;
860 for c in text.chars() {
861 if chars >= MAX_CONTENT {
862 break;
863 }
864 if c.is_whitespace() {
865 if !last_was_space {
866 out.push(' ');
867 chars += 1;
868 last_was_space = true;
869 }
870 } else {
871 out.push(c);
872 chars += 1;
873 last_was_space = false;
874 }
875 }
876 out.trim().to_owned()
877}
878
879#[derive(Debug, Clone, Copy, Default)]
883pub struct FileNodeExtractor;
884
885impl Extractor for FileNodeExtractor {
886 fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
887 FactSet::new().with_node(file_node(
888 path,
889 blob_id,
890 bytes,
891 None,
892 IngestConfig::default(),
893 ))
894 }
895}
896
897#[derive(Debug, Clone, Copy, Default)]
904pub struct RustExtractor;
905
906impl Extractor for RustExtractor {
907 fn extract(&self, path: &str, blob_id: &str, bytes: &[u8]) -> FactSet {
908 rust_facts(path, blob_id, bytes, IngestConfig::default())
909 }
910}
911
912fn rust_facts(path: &str, blob_id: &str, bytes: &[u8], ingest: IngestConfig) -> FactSet {
915 let mut parser = tree_sitter::Parser::new();
916 if parser
919 .set_language(&tree_sitter_rust::LANGUAGE.into())
920 .is_err()
921 {
922 return FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
923 }
924 let Some(tree) = parser.parse(bytes, None) else {
925 return FactSet::new().with_node(file_node(path, blob_id, bytes, None, ingest));
926 };
927
928 let mut walk = RustWalk {
929 path,
930 blob_id,
931 src: bytes,
932 nodes: vec![file_node(path, blob_id, bytes, Some("rust"), ingest)],
933 edges: Vec::new(),
934 };
935 let root = tree.root_node();
936 let mut cursor = root.walk();
937 let children: Vec<_> = root.children(&mut cursor).collect();
938 for child in children {
939 walk.visit(child, &[]);
940 }
941 walk.synthesize_config_keys(root);
946
947 walk.nodes.sort_by(|a, b| a.key.cmp(&b.key));
950 walk.edges
951 .sort_by(|a, b| (a.kind.as_str(), &a.src, &a.dst).cmp(&(b.kind.as_str(), &b.src, &b.dst)));
952 FactSet {
953 nodes: walk.nodes,
954 edges: walk.edges,
955 }
956}
957
958struct Scope {
962 seg: String,
963 key: Option<String>,
964}
965
966struct FieldDef {
969 name: String,
970 type_idents: Vec<String>,
971}
972
973const TRANSPARENT_WRAPPERS: &[&str] = &[
977 "Option", "Box", "Arc", "Rc", "Cow", "RefCell", "Cell", "Mutex", "RwLock",
978];
979
980const COLLECTION_WRAPPERS: &[&str] = &[
985 "Vec", "VecDeque", "HashMap", "BTreeMap", "HashSet", "BTreeSet", "IndexMap",
986];
987
988fn core_type_name(type_idents: &[String]) -> Option<String> {
993 type_idents
994 .iter()
995 .find(|t| !TRANSPARENT_WRAPPERS.contains(&t.as_str()))
996 .or_else(|| type_idents.first())
997 .cloned()
998}
999
1000fn recursion_target<'a>(
1005 type_idents: &'a [String],
1006 known: &std::collections::BTreeMap<String, StructDef>,
1007) -> Option<&'a str> {
1008 for t in type_idents {
1009 if COLLECTION_WRAPPERS.contains(&t.as_str()) {
1010 return None;
1011 }
1012 if TRANSPARENT_WRAPPERS.contains(&t.as_str()) {
1013 continue;
1014 }
1015 return known.contains_key(t).then_some(t.as_str());
1016 }
1017 None
1018}
1019
1020struct StructDef {
1023 fields: Vec<FieldDef>,
1024 is_root: bool,
1025}
1026
1027const MAX_CONFIG_DEPTH: usize = 16;
1029
1030fn expand_config_keys(
1037 table: &std::collections::BTreeMap<String, StructDef>,
1038 struct_name: &str,
1039 prefix: &str,
1040 root: &str,
1041 visited: &mut std::collections::BTreeSet<String>,
1042 depth: usize,
1043 out: &mut std::collections::BTreeMap<String, String>,
1044) {
1045 let Some(def) = table.get(struct_name) else {
1046 return;
1047 };
1048 for f in &def.fields {
1049 let key = if prefix.is_empty() {
1050 f.name.clone()
1051 } else {
1052 format!("{prefix}.{}", f.name)
1053 };
1054 match recursion_target(&f.type_idents, table) {
1055 Some(inner) if depth < MAX_CONFIG_DEPTH && !visited.contains(inner) => {
1056 visited.insert(inner.to_owned());
1057 expand_config_keys(table, inner, &key, root, visited, depth + 1, out);
1058 visited.remove(inner);
1059 }
1060 _ => {
1061 out.entry(key).or_insert_with(|| root.to_owned());
1062 }
1063 }
1064 }
1065}
1066
1067struct RustWalk<'a> {
1069 path: &'a str,
1070 blob_id: &'a str,
1071 src: &'a [u8],
1072 nodes: Vec<Node>,
1073 edges: Vec<Edge>,
1074}
1075
1076impl RustWalk<'_> {
1077 fn visit(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
1079 match node.kind() {
1080 "function_item" => self.visit_symbol(node, scope, NodeKind::Fn, true),
1081 "struct_item" | "union_item" => self.visit_symbol(node, scope, NodeKind::Struct, false),
1082 "enum_item" => self.visit_symbol(node, scope, NodeKind::Enum, false),
1083 "trait_item" => self.visit_symbol(node, scope, NodeKind::Trait, false),
1084 "mod_item" => self.visit_symbol(node, scope, NodeKind::Module, false),
1085 "type_item" => self.visit_symbol(node, scope, NodeKind::Other("type".into()), false),
1086 "macro_definition" => {
1087 self.visit_symbol(node, scope, NodeKind::Other("macro".into()), false);
1088 }
1089 "impl_item" => self.visit_impl(node, scope),
1090 "use_declaration" => self.visit_use(node),
1091 _ => self.visit_children(node, scope),
1094 }
1095 }
1096
1097 fn visit_children(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
1099 let mut cursor = node.walk();
1100 let children: Vec<_> = node.named_children(&mut cursor).collect();
1101 for child in children {
1102 self.visit(child, scope);
1103 }
1104 }
1105
1106 fn visit_symbol(
1109 &mut self,
1110 node: tree_sitter::Node,
1111 scope: &[Scope],
1112 kind: NodeKind,
1113 collect_calls: bool,
1114 ) {
1115 let Some(name) = self.field_text(node, "name") else {
1116 return self.visit_children(node, scope);
1117 };
1118 let qualified = qualify(scope, &name);
1119 let key = format!("sym:rust:{}#{qualified}", self.path);
1120
1121 let mut meta = serde_json::Map::new();
1122 if collect_calls {
1123 let mut calls = Vec::new();
1124 self.collect_calls(node, &mut calls);
1125 calls.sort();
1126 calls.dedup();
1127 if !calls.is_empty() {
1128 meta.insert("calls".into(), serde_json::Value::from(calls));
1129 }
1130 }
1131 if let Some(doc) = self.doc_comment(node) {
1133 meta.insert("content".into(), serde_json::Value::from(doc));
1134 }
1135 if matches!(node.kind(), "struct_item" | "union_item") {
1147 let defs = self.struct_fields(node);
1148 if !defs.is_empty() {
1149 let names: Vec<&str> = defs.iter().map(|f| f.name.as_str()).collect();
1150 meta.insert("fields".into(), serde_json::Value::from(names));
1151 let types: serde_json::Map<String, serde_json::Value> = defs
1152 .iter()
1153 .filter_map(|f| {
1154 core_type_name(&f.type_idents).map(|t| (f.name.clone(), t.into()))
1155 })
1156 .collect();
1157 if !types.is_empty() {
1158 meta.insert("field_types".into(), serde_json::Value::Object(types));
1159 }
1160 }
1161 if self.has_config_marker(node) {
1162 meta.insert("config_root".into(), serde_json::Value::Bool(true));
1163 }
1164 }
1165
1166 self.nodes.push(Node {
1167 key: key.clone(),
1168 kind,
1169 name,
1170 path: Some(self.path.to_owned()),
1171 lang: Some("rust".to_owned()),
1172 blob_hash: Some(self.blob_id.to_owned()),
1173 span: Some(span(node)),
1174 provenance: Provenance::Derived,
1175 meta: serde_json::Value::Object(meta),
1176 });
1177 self.link_parent(&key, scope);
1178
1179 let child_scope = extend(scope, &self.simple(node, "name"), Some(key));
1182 self.recurse_body(node, &child_scope);
1183 }
1184
1185 fn doc_comment(&self, node: tree_sitter::Node) -> Option<String> {
1189 let mut parts: Vec<String> = Vec::new();
1190 let mut prev = node.prev_sibling();
1191 while let Some(n) = prev {
1192 match n.kind() {
1193 "line_comment" | "block_comment" => match doc_comment_body(self.text(n)) {
1194 Some(body) => {
1195 parts.push(body);
1196 prev = n.prev_sibling();
1197 }
1198 None => break,
1199 },
1200 "attribute_item" => prev = n.prev_sibling(),
1201 _ => break,
1202 }
1203 }
1204 if parts.is_empty() {
1205 return None;
1206 }
1207 parts.reverse();
1208 let joined = cap_content(&parts.join(" "));
1209 (!joined.is_empty()).then_some(joined)
1210 }
1211
1212 fn visit_impl(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
1215 let type_name = self
1216 .field_text(node, "type")
1217 .unwrap_or_else(|| "impl".to_owned());
1218 let child_scope = extend(scope, &type_name, None);
1219 self.recurse_body(node, &child_scope);
1220 }
1221
1222 fn visit_use(&mut self, node: tree_sitter::Node) {
1225 let Some(arg) = node.child_by_field_name("argument") else {
1226 return;
1227 };
1228 let text: String = self
1229 .text(arg)
1230 .chars()
1231 .filter(|c| !c.is_whitespace())
1232 .collect();
1233 if text.is_empty() {
1234 return;
1235 }
1236 let key = format!("import:rust:{text}");
1237 self.nodes.push(Node {
1238 key: key.clone(),
1239 kind: NodeKind::Other("import".into()),
1240 name: text,
1241 path: None,
1242 lang: Some("rust".to_owned()),
1243 blob_hash: None,
1244 span: None,
1245 provenance: Provenance::Derived,
1246 meta: serde_json::Value::Null,
1247 });
1248 self.edges
1249 .push(Edge::derived(file_key(self.path), key, EdgeKind::Imports));
1250 }
1251
1252 fn link_parent(&mut self, key: &str, scope: &[Scope]) {
1255 if let Some(parent) = scope.iter().rev().find_map(|s| s.key.as_deref()) {
1256 self.edges.push(Edge::derived(
1257 parent.to_owned(),
1258 key.to_owned(),
1259 EdgeKind::Contains,
1260 ));
1261 } else {
1262 self.edges.push(Edge::derived(
1263 file_key(self.path),
1264 key.to_owned(),
1265 EdgeKind::Defines,
1266 ));
1267 }
1268 }
1269
1270 fn struct_fields(&self, node: tree_sitter::Node) -> Vec<FieldDef> {
1277 let mut out = Vec::new();
1278 let mut cursor = node.walk();
1279 for child in node.named_children(&mut cursor) {
1280 if child.kind() == "field_declaration_list" {
1281 let mut inner = child.walk();
1282 for field in child.named_children(&mut inner) {
1283 if field.kind() == "field_declaration"
1284 && let Some(name) = field.child_by_field_name("name")
1285 {
1286 let type_idents = field
1287 .child_by_field_name("type")
1288 .map(|t| self.type_idents(t))
1289 .unwrap_or_default();
1290 out.push(FieldDef {
1291 name: self.text(name).to_owned(),
1292 type_idents,
1293 });
1294 }
1295 }
1296 }
1297 }
1298 out
1299 }
1300
1301 fn type_idents(&self, ty: tree_sitter::Node) -> Vec<String> {
1306 let mut out = Vec::new();
1307 self.collect_type_idents(ty, &mut out);
1308 out
1309 }
1310
1311 fn collect_type_idents(&self, node: tree_sitter::Node, out: &mut Vec<String>) {
1312 if matches!(node.kind(), "type_identifier" | "primitive_type") {
1316 out.push(self.text(node).to_owned());
1317 }
1318 let mut cursor = node.walk();
1319 for child in node.named_children(&mut cursor) {
1320 self.collect_type_idents(child, out);
1321 }
1322 }
1323
1324 fn has_config_marker(&self, node: tree_sitter::Node) -> bool {
1335 const MARKER: &str = "@rto:config";
1336 let mut prev = node.prev_sibling();
1337 while let Some(n) = prev {
1338 match n.kind() {
1339 "line_comment" | "block_comment" | "attribute_item" => {
1340 if self.text(n).contains(MARKER) {
1341 return true;
1342 }
1343 prev = n.prev_sibling();
1344 }
1345 _ => break,
1346 }
1347 }
1348 false
1349 }
1350
1351 fn recurse_body(&mut self, node: tree_sitter::Node, scope: &[Scope]) {
1353 let mut cursor = node.walk();
1354 let children: Vec<_> = node.named_children(&mut cursor).collect();
1355 for child in children {
1356 match child.kind() {
1357 "declaration_list" | "field_declaration_list" | "trait_body" => {
1358 self.visit_children(child, scope);
1359 }
1360 _ => {}
1361 }
1362 }
1363 }
1364
1365 fn collect_calls(&self, node: tree_sitter::Node, out: &mut Vec<String>) {
1368 let mut cursor = node.walk();
1369 for child in node.named_children(&mut cursor) {
1370 if child.kind() == "call_expression"
1371 && let Some(func) = child.child_by_field_name("function")
1372 && let Some(name) = self.callee_name(func)
1373 {
1374 out.push(name);
1375 }
1376 self.collect_calls(child, out);
1377 }
1378 }
1379
1380 fn callee_name(&self, func: tree_sitter::Node) -> Option<String> {
1391 match func.kind() {
1392 "identifier" => Some(self.text(func).to_owned()),
1393 "scoped_identifier" => {
1394 let name = func.child_by_field_name("name")?;
1395 let qualifier = func
1398 .child_by_field_name("path")
1399 .and_then(|p| self.text(p).rsplit("::").next().map(str::to_owned));
1400 Some(qualify_callee(qualifier.as_deref(), self.text(name)))
1401 }
1402 "field_expression" => {
1403 let name = func.child_by_field_name("field")?;
1404 let on_self = func
1407 .child_by_field_name("value")
1408 .is_some_and(|v| self.text(v) == "self");
1409 Some(qualify_callee(on_self.then_some("Self"), self.text(name)))
1410 }
1411 _ => None,
1412 }
1413 }
1414
1415 fn synthesize_config_keys(&mut self, root: tree_sitter::Node) {
1439 let table = self.collect_struct_defs(root);
1440 let mut keys: std::collections::BTreeMap<String, String> =
1443 std::collections::BTreeMap::new();
1444 for (name, def) in &table {
1445 if !def.is_root {
1446 continue;
1447 }
1448 let mut visited = std::collections::BTreeSet::new();
1449 visited.insert(name.clone());
1450 expand_config_keys(&table, name, "", name, &mut visited, 0, &mut keys);
1451 }
1452 let file = file_key(self.path);
1453 for (dotted, root_name) in keys {
1454 let node_key = format!("cfgkey:{}#{dotted}", self.path);
1455 let mut node = Node::new(
1456 node_key.clone(),
1457 NodeKind::Other(crate::config_keys::KIND.into()),
1458 dotted.clone(),
1459 );
1460 node.path = Some(self.path.to_owned());
1461 node.blob_hash = Some(self.blob_id.to_owned());
1462 node.meta = serde_json::json!({
1469 "key": dotted,
1470 "source": "struct",
1471 "struct": root_name,
1472 });
1473 self.edges.push(Edge::derived(
1474 file.clone(),
1475 node_key.clone(),
1476 EdgeKind::Contains,
1477 ));
1478 self.nodes.push(node);
1479 }
1480 }
1481
1482 fn collect_struct_defs(
1486 &self,
1487 root: tree_sitter::Node,
1488 ) -> std::collections::BTreeMap<String, StructDef> {
1489 let mut out = std::collections::BTreeMap::new();
1490 self.collect_struct_defs_into(root, &mut out);
1491 out
1492 }
1493
1494 fn collect_struct_defs_into(
1495 &self,
1496 node: tree_sitter::Node,
1497 out: &mut std::collections::BTreeMap<String, StructDef>,
1498 ) {
1499 if matches!(node.kind(), "struct_item" | "union_item")
1500 && let Some(name) = self.field_text(node, "name")
1501 {
1502 out.entry(name.clone()).or_insert_with(|| StructDef {
1503 fields: self.struct_fields(node),
1504 is_root: self.has_config_marker(node),
1505 });
1506 }
1507 let mut cursor = node.walk();
1508 for child in node.named_children(&mut cursor) {
1509 self.collect_struct_defs_into(child, out);
1510 }
1511 }
1512
1513 fn text(&self, node: tree_sitter::Node) -> &str {
1514 node.utf8_text(self.src).unwrap_or("")
1515 }
1516
1517 fn field_text(&self, node: tree_sitter::Node, field: &str) -> Option<String> {
1518 node.child_by_field_name(field)
1519 .map(|n| self.text(n).to_owned())
1520 }
1521
1522 fn simple(&self, node: tree_sitter::Node, field: &str) -> String {
1523 self.field_text(node, field).unwrap_or_default()
1524 }
1525}
1526
1527struct TagLang {
1543 lang: &'static str,
1545 grammar_key: &'static str,
1550 language: tree_sitter::Language,
1552 query: std::borrow::Cow<'static, str>,
1556}
1557
1558#[allow(clippy::too_many_lines)]
1563fn tag_lang_for(ext: &str) -> Option<TagLang> {
1564 use std::borrow::Cow;
1565 let ts_query = || -> Cow<'static, str> {
1569 Cow::Owned(format!(
1570 "{}\n{}",
1571 tree_sitter_javascript::TAGS_QUERY,
1572 tree_sitter_typescript::TAGS_QUERY
1573 ))
1574 };
1575 let borrowed = |q: &'static str| -> Cow<'static, str> { Cow::Borrowed(q) };
1576
1577 let (lang, language, query): (&str, tree_sitter::Language, Cow<'static, str>) = match ext {
1578 "py" | "pyi" => (
1579 "python",
1580 tree_sitter_python::LANGUAGE.into(),
1581 borrowed(tree_sitter_python::TAGS_QUERY),
1582 ),
1583 "js" | "jsx" | "mjs" | "cjs" => (
1584 "javascript",
1585 tree_sitter_javascript::LANGUAGE.into(),
1586 borrowed(tree_sitter_javascript::TAGS_QUERY),
1587 ),
1588 "ts" | "mts" | "cts" => (
1589 "typescript",
1590 tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
1591 ts_query(),
1592 ),
1593 "tsx" => (
1594 "tsx",
1595 tree_sitter_typescript::LANGUAGE_TSX.into(),
1596 ts_query(),
1597 ),
1598 "go" => (
1599 "go",
1600 tree_sitter_go::LANGUAGE.into(),
1601 borrowed(tree_sitter_go::TAGS_QUERY),
1602 ),
1603 "rb" => (
1604 "ruby",
1605 tree_sitter_ruby::LANGUAGE.into(),
1606 borrowed(tree_sitter_ruby::TAGS_QUERY),
1607 ),
1608 "java" => (
1609 "java",
1610 tree_sitter_java::LANGUAGE.into(),
1611 borrowed(tree_sitter_java::TAGS_QUERY),
1612 ),
1613 "c" | "h" => (
1614 "c",
1615 tree_sitter_c::LANGUAGE.into(),
1616 borrowed(tree_sitter_c::TAGS_QUERY),
1617 ),
1618 "cc" | "cpp" | "cxx" | "hpp" | "hh" | "hxx" => (
1619 "cpp",
1620 tree_sitter_cpp::LANGUAGE.into(),
1621 borrowed(tree_sitter_cpp::TAGS_QUERY),
1622 ),
1623 "cs" => (
1626 "csharp",
1627 tree_sitter_c_sharp::LANGUAGE.into(),
1628 borrowed(include_str!("queries/csharp/tags.scm")),
1629 ),
1630 "php" => (
1631 "php",
1632 tree_sitter_php::LANGUAGE_PHP.into(),
1633 borrowed(tree_sitter_php::TAGS_QUERY),
1634 ),
1635 "scala" | "sc" => (
1637 "scala",
1638 tree_sitter_scala::LANGUAGE.into(),
1639 borrowed(include_str!("queries/scala/tags.scm")),
1640 ),
1641 "ml" => (
1642 "ocaml",
1643 tree_sitter_ocaml::LANGUAGE_OCAML.into(),
1644 borrowed(tree_sitter_ocaml::TAGS_QUERY),
1645 ),
1646 "mli" => (
1647 "ocaml",
1648 tree_sitter_ocaml::LANGUAGE_OCAML_INTERFACE.into(),
1649 borrowed(tree_sitter_ocaml::TAGS_QUERY),
1650 ),
1651 "ex" | "exs" => (
1652 "elixir",
1653 tree_sitter_elixir::LANGUAGE.into(),
1654 borrowed(tree_sitter_elixir::TAGS_QUERY),
1655 ),
1656 "sh" | "bash" => (
1658 "bash",
1659 tree_sitter_bash::LANGUAGE.into(),
1660 borrowed(include_str!("queries/bash/tags.scm")),
1661 ),
1662 "sql" => (
1664 "sql",
1665 tree_sitter_sequel::LANGUAGE.into(),
1666 borrowed(include_str!("queries/sql/tags.scm")),
1667 ),
1668 _ => return None,
1669 };
1670 let grammar_key = match ext {
1673 "mli" => "ocaml-interface",
1674 _ => lang,
1675 };
1676 Some(TagLang {
1677 lang,
1678 grammar_key,
1679 language,
1680 query,
1681 })
1682}
1683
1684type TagConfig = std::sync::Arc<tree_sitter_tags::TagsConfiguration>;
1686
1687static TAG_CONFIGS: std::sync::LazyLock<
1694 std::sync::Mutex<std::collections::HashMap<&'static str, Option<TagConfig>>>,
1695> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
1696
1697fn tag_config(def: &TagLang) -> Option<TagConfig> {
1700 let mut cache = TAG_CONFIGS
1701 .lock()
1702 .unwrap_or_else(std::sync::PoisonError::into_inner);
1703 cache
1704 .entry(def.grammar_key)
1705 .or_insert_with(|| {
1706 tree_sitter_tags::TagsConfiguration::new(def.language.clone(), &def.query, "")
1707 .ok()
1708 .map(std::sync::Arc::new)
1709 })
1710 .clone()
1711}
1712
1713fn import_query_for(lang: &str) -> Option<&'static str> {
1721 Some(match lang {
1722 "python" => {
1724 "(import_statement name: (dotted_name) @path)\n\
1725 (import_statement name: (aliased_import name: (dotted_name) @path))\n\
1726 (import_from_statement module_name: (dotted_name) @path)\n\
1727 (import_from_statement module_name: (relative_import) @path)"
1728 }
1729 "javascript" | "typescript" | "tsx" => {
1731 "(import_statement source: (string (string_fragment) @path))\n\
1732 (export_statement source: (string (string_fragment) @path))"
1733 }
1734 "go" => "(import_spec path: (interpreted_string_literal) @path)",
1736 "java" => {
1738 "(import_declaration (scoped_identifier) @path)\n\
1739 (import_declaration (identifier) @path)"
1740 }
1741 "c" | "cpp" => {
1743 "(preproc_include path: (string_literal) @path)\n\
1744 (preproc_include path: (system_lib_string) @path)"
1745 }
1746 _ => return None,
1747 })
1748}
1749
1750type ImportQuery = std::sync::Arc<tree_sitter::Query>;
1752
1753static IMPORT_QUERIES: std::sync::LazyLock<
1757 std::sync::Mutex<std::collections::HashMap<&'static str, Option<ImportQuery>>>,
1758> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
1759
1760fn import_query(def: &TagLang) -> Option<ImportQuery> {
1762 let mut cache = IMPORT_QUERIES
1763 .lock()
1764 .unwrap_or_else(std::sync::PoisonError::into_inner);
1765 cache
1766 .entry(def.grammar_key)
1767 .or_insert_with(|| {
1768 let src = import_query_for(def.lang)?;
1769 tree_sitter::Query::new(&def.language, src)
1770 .ok()
1771 .map(std::sync::Arc::new)
1772 })
1773 .clone()
1774}
1775
1776fn normalize_import(raw: &str) -> String {
1779 raw.trim()
1780 .trim_matches(|c| c == '"' || c == '\'' || c == '<' || c == '>')
1781 .trim()
1782 .to_owned()
1783}
1784
1785fn append_import_facts(
1789 path: &str,
1790 def: &TagLang,
1791 bytes: &[u8],
1792 nodes: &mut Vec<Node>,
1793 edges: &mut Vec<Edge>,
1794) {
1795 use streaming_iterator::StreamingIterator as _;
1796
1797 let Some(query) = import_query(def) else {
1798 return;
1799 };
1800 let mut parser = tree_sitter::Parser::new();
1801 if parser.set_language(&def.language).is_err() {
1802 return;
1803 }
1804 let Some(tree) = parser.parse(bytes, None) else {
1805 return;
1806 };
1807 let mut cursor = tree_sitter::QueryCursor::new();
1808 let mut seen = std::collections::BTreeSet::new();
1809 let mut matches = cursor.matches(&query, tree.root_node(), bytes);
1810 while let Some(m) = matches.next() {
1811 for cap in m.captures {
1812 let Ok(raw) = cap.node.utf8_text(bytes) else {
1813 continue;
1814 };
1815 let module = normalize_import(raw);
1816 if module.is_empty() {
1817 continue;
1818 }
1819 let key = format!("import:{}:{module}", def.lang);
1820 if seen.insert(key.clone()) {
1821 nodes.push(Node {
1822 key: key.clone(),
1823 kind: NodeKind::Other("import".into()),
1824 name: module,
1825 path: None,
1829 lang: Some(def.lang.to_owned()),
1830 blob_hash: None,
1831 span: None,
1832 provenance: Provenance::Derived,
1833 meta: serde_json::Value::Null,
1834 });
1835 edges.push(Edge::derived(file_key(path), key, EdgeKind::Imports));
1836 }
1837 }
1838 }
1839}
1840
1841fn tag_node_kind(syntax_type: &str) -> NodeKind {
1844 match syntax_type {
1845 "function" | "method" | "constructor" => NodeKind::Fn,
1846 "class" | "struct" => NodeKind::Struct,
1847 "interface" | "trait" | "protocol" => NodeKind::Trait,
1848 "enum" => NodeKind::Enum,
1849 "module" | "namespace" | "object" => NodeKind::Module,
1851 other => NodeKind::Other(other.to_owned()),
1852 }
1853}
1854
1855struct TagDef {
1857 name: String,
1858 kind: NodeKind,
1859 range: std::ops::Range<usize>,
1860 docs: Option<String>,
1861}
1862
1863fn tag_facts(
1867 path: &str,
1868 blob_id: &str,
1869 bytes: &[u8],
1870 ext: &str,
1871 ingest: IngestConfig,
1872) -> Option<FactSet> {
1873 let def = tag_lang_for(ext)?;
1874 let lang = def.lang;
1875 let config = tag_config(&def)?;
1876
1877 let mut ctx = tree_sitter_tags::TagsContext::new();
1878 let (tags, _had_error) = ctx.generate_tags(&config, bytes, None).ok()?;
1879
1880 let mut defs: Vec<TagDef> = Vec::new();
1881 let mut calls: Vec<(usize, String)> = Vec::new();
1884 for tag in tags {
1885 let Ok(tag) = tag else { continue };
1886 let Some(name) = bytes
1887 .get(tag.name_range.clone())
1888 .and_then(|b| std::str::from_utf8(b).ok())
1889 else {
1890 continue;
1891 };
1892 let syntax = config.syntax_type_name(tag.syntax_type_id);
1893 if tag.is_definition {
1894 defs.push(TagDef {
1895 name: name.to_owned(),
1896 kind: tag_node_kind(syntax),
1897 range: tag.range.clone(),
1898 docs: tag.docs.clone(),
1900 });
1901 } else if syntax == "call" || syntax == "send" {
1902 calls.push((tag.range.start, name.to_owned()));
1904 }
1905 }
1906
1907 let parents: Vec<Option<usize>> = (0..defs.len())
1912 .map(|i| smallest_enclosing(&defs, defs[i].range.clone(), Some(i)))
1913 .collect();
1914
1915 let keys: Vec<String> = (0..defs.len())
1916 .map(|i| {
1917 let qualified = qualified_name(&defs, &parents, i);
1918 format!("sym:{lang}:{path}#{qualified}")
1919 })
1920 .collect();
1921
1922 let mut nodes = vec![file_node(path, blob_id, bytes, Some(lang), ingest)];
1923 let mut edges: Vec<Edge> = Vec::new();
1924
1925 for (i, d) in defs.iter().enumerate() {
1926 let mut meta = serde_json::Map::new();
1927 if let Some(doc) = &d.docs {
1928 let content = cap_content(doc);
1929 if !content.is_empty() {
1930 meta.insert("content".into(), serde_json::Value::from(content));
1931 }
1932 }
1933 if d.kind == NodeKind::Fn {
1936 let mut names: Vec<String> = calls
1937 .iter()
1938 .filter(|(off, _)| d.range.contains(off))
1939 .filter(|(off, _)| smallest_enclosing_off(&defs, *off) == Some(i))
1940 .map(|(_, name)| name.clone())
1941 .collect();
1942 names.sort();
1943 names.dedup();
1944 if !names.is_empty() {
1945 meta.insert("calls".into(), serde_json::Value::from(names));
1946 }
1947 }
1948
1949 let start = u32::try_from(d.range.start).unwrap_or(u32::MAX);
1950 let end = u32::try_from(d.range.end).unwrap_or(u32::MAX);
1951 nodes.push(Node {
1952 key: keys[i].clone(),
1953 kind: d.kind.clone(),
1954 name: d.name.clone(),
1955 path: Some(path.to_owned()),
1956 lang: Some(lang.to_owned()),
1957 blob_hash: Some(blob_id.to_owned()),
1958 span: Some(Span::new(start, end)),
1959 provenance: Provenance::Derived,
1960 meta: serde_json::Value::Object(meta),
1961 });
1962
1963 match parents[i] {
1964 Some(p) => edges.push(Edge::derived(
1965 keys[p].clone(),
1966 keys[i].clone(),
1967 EdgeKind::Contains,
1968 )),
1969 None => edges.push(Edge::derived(
1970 file_key(path),
1971 keys[i].clone(),
1972 EdgeKind::Defines,
1973 )),
1974 }
1975 }
1976
1977 append_import_facts(path, &def, bytes, &mut nodes, &mut edges);
1979
1980 nodes.sort_by(|a, b| a.key.cmp(&b.key));
1983 nodes.dedup_by(|a, b| a.key == b.key);
1984 edges.sort_by(|a, b| (a.kind.as_str(), &a.src, &a.dst).cmp(&(b.kind.as_str(), &b.src, &b.dst)));
1985 edges.dedup();
1986 Some(FactSet { nodes, edges })
1987}
1988
1989fn smallest_enclosing(
1992 defs: &[TagDef],
1993 range: std::ops::Range<usize>,
1994 skip: Option<usize>,
1995) -> Option<usize> {
1996 let mut best: Option<usize> = None;
1997 for (j, c) in defs.iter().enumerate() {
1998 if Some(j) == skip {
1999 continue;
2000 }
2001 let encloses = c.range.start <= range.start
2003 && c.range.end >= range.end
2004 && (c.range.end - c.range.start) > (range.end - range.start);
2005 if encloses
2006 && best.is_none_or(|b| {
2007 defs[b].range.end - defs[b].range.start > c.range.end - c.range.start
2008 })
2009 {
2010 best = Some(j);
2011 }
2012 }
2013 best
2014}
2015
2016fn smallest_enclosing_off(defs: &[TagDef], off: usize) -> Option<usize> {
2018 let mut best: Option<usize> = None;
2019 for (j, c) in defs.iter().enumerate() {
2020 if c.range.contains(&off)
2021 && best.is_none_or(|b| {
2022 defs[b].range.end - defs[b].range.start > c.range.end - c.range.start
2023 })
2024 {
2025 best = Some(j);
2026 }
2027 }
2028 best
2029}
2030
2031fn qualified_name(defs: &[TagDef], parents: &[Option<usize>], i: usize) -> String {
2034 let mut chain: Vec<&str> = vec![defs[i].name.as_str()];
2035 let mut cur = parents[i];
2036 let mut guard = defs.len();
2039 while let Some(p) = cur {
2040 if guard == 0 {
2041 break;
2042 }
2043 guard -= 1;
2044 chain.push(defs[p].name.as_str());
2045 cur = parents[p];
2046 }
2047 chain.reverse();
2048 chain.join("::")
2049}
2050
2051fn span(node: tree_sitter::Node) -> Span {
2053 let start = u32::try_from(node.start_byte()).unwrap_or(u32::MAX);
2054 let end = u32::try_from(node.end_byte()).unwrap_or(u32::MAX);
2055 Span::new(start, end)
2056}
2057
2058fn qualify(scope: &[Scope], name: &str) -> String {
2060 let mut parts: Vec<&str> = scope.iter().map(|s| s.seg.as_str()).collect();
2061 parts.push(name);
2062 parts.join("::")
2063}
2064
2065fn qualify_callee(qualifier: Option<&str>, name: &str) -> String {
2070 match qualifier {
2071 Some(q) if !q.is_empty() && !matches!(q, "self" | "crate" | "super") => {
2072 format!("{q}::{name}")
2073 }
2074 _ => name.to_owned(),
2075 }
2076}
2077
2078fn extend(scope: &[Scope], seg: &str, key: Option<String>) -> Vec<Scope> {
2080 let mut next: Vec<Scope> = scope
2081 .iter()
2082 .map(|s| Scope {
2083 seg: s.seg.clone(),
2084 key: s.key.clone(),
2085 })
2086 .collect();
2087 next.push(Scope {
2088 seg: seg.to_owned(),
2089 key,
2090 });
2091 next
2092}
2093
2094#[cfg(test)]
2095mod tests {
2096 use super::{Extractor, FileNodeExtractor, Registry, RustExtractor};
2097 use crate::{EdgeKind, Node, NodeKind};
2098
2099 #[test]
2100 fn file_node_extractor_is_deterministic_and_tagged() {
2101 let ex = FileNodeExtractor;
2102 let a = ex.extract("src/lib.rs", "abc123", b"one\ntwo\n");
2103 let b = ex.extract("src/lib.rs", "abc123", b"one\ntwo\n");
2104 assert_eq!(a, b, "extraction must be deterministic");
2105
2106 assert_eq!(a.nodes.len(), 1);
2107 assert!(a.edges.is_empty());
2108 let node = &a.nodes[0];
2109 assert_eq!(node.key, "file:src/lib.rs");
2110 assert_eq!(node.kind, NodeKind::File);
2111 assert_eq!(node.name, "lib.rs");
2112 assert_eq!(node.blob_hash.as_deref(), Some("abc123"));
2113 assert_eq!(node.meta["lines"], 2);
2114 assert_eq!(node.meta["bytes"], 8);
2115 }
2116
2117 #[test]
2118 fn config_files_emit_config_key_nodes() {
2119 let reg = Registry::new(crate::IngestConfig::default());
2120 let toml = b"[serve]\naddr = \"0.0.0.0:8443\"\ntools = false\n";
2121 let a = reg.extract("config.toml", "cfg1", toml);
2122 let b = reg.extract("config.toml", "cfg1", toml);
2123 assert_eq!(a, b, "config extraction must be deterministic");
2124
2125 assert!(a.nodes.iter().any(|n| n.key == "file:config.toml"));
2127 let addr = a
2128 .nodes
2129 .iter()
2130 .find(|n| n.key == "cfgkey:config.toml#serve.addr")
2131 .expect("serve.addr config_key node");
2132 assert_eq!(addr.kind, NodeKind::Other("config_key".into()));
2133 assert_eq!(addr.name, "serve.addr");
2134 assert_eq!(addr.meta["value"], "0.0.0.0:8443"); assert!(a.edges.iter().any(|e| {
2137 e.src == "file:config.toml"
2138 && e.dst == "cfgkey:config.toml#serve.addr"
2139 && e.kind == EdgeKind::Contains
2140 }));
2141
2142 let env = reg.extract(".env", "env1", b"PORT=8080\nPORT=9090\nAPI_TOKEN=s3cr3t\n");
2145 let port = env
2146 .nodes
2147 .iter()
2148 .find(|n| n.key == "cfgkey:.env#PORT")
2149 .expect("PORT node");
2150 assert_eq!(port.meta["value"], "9090", "dotenv last-one-wins");
2151 assert_eq!(
2152 env.nodes
2153 .iter()
2154 .filter(|n| n.key == "cfgkey:.env#PORT")
2155 .count(),
2156 1
2157 );
2158 let token = env
2159 .nodes
2160 .iter()
2161 .find(|n| n.key == "cfgkey:.env#API_TOKEN")
2162 .expect("API_TOKEN node");
2163 assert_eq!(token.meta["value"], "<redacted>", "secret not persisted");
2164 let rs = reg.extract("src/lib.rs", "x", b"pub fn f() {}\n");
2166 assert!(
2167 rs.nodes
2168 .iter()
2169 .all(|n| n.kind != NodeKind::Other("config_key".into()))
2170 );
2171 }
2172
2173 #[test]
2174 fn dockerfile_emits_image_ref_nodes_and_skips_internal_stages() {
2175 let reg = Registry::new(crate::IngestConfig::default());
2176 let df = b"FROM --platform=linux/amd64 rust:1.90 AS builder\nRUN cargo build\n\
2179 FROM builder AS test\nFROM registry.io/app:1.2@sha256:abc AS run\nFROM scratch\n";
2180 let a = reg.extract("Dockerfile", "d1", df);
2181 let b = reg.extract("Dockerfile", "d1", df);
2182 assert_eq!(a, b, "dockerfile extraction must be deterministic");
2183
2184 let refs: Vec<&Node> = a
2185 .nodes
2186 .iter()
2187 .filter(|n| n.kind == NodeKind::Other("image_ref".into()))
2188 .collect();
2189 assert_eq!(refs.len(), 2, "got: {refs:?}");
2192 let rust = refs
2193 .iter()
2194 .find(|n| n.meta["image"] == "rust")
2195 .expect("rust");
2196 assert_eq!(rust.meta["tag"], "1.90");
2197 let app = refs
2198 .iter()
2199 .find(|n| n.meta["image"] == "registry.io/app:1.2")
2200 .expect("app digest");
2201 assert_eq!(app.meta["digest"], "sha256:abc");
2202 assert!(
2204 a.edges
2205 .iter()
2206 .any(|e| { e.src == "file:Dockerfile" && e.kind == EdgeKind::References })
2207 );
2208 assert!(
2210 reg.extract("Dockerfile.prod", "d2", b"FROM alpine:3\n")
2211 .nodes
2212 .iter()
2213 .any(|n| n.kind == NodeKind::Other("image_ref".into()))
2214 );
2215
2216 let c = reg.extract("Dockerfile", "d3", b"FROM alpine AS alpine\n");
2219 assert!(
2220 c.nodes
2221 .iter()
2222 .any(|n| n.kind == NodeKind::Other("image_ref".into())
2223 && n.meta["image"] == "alpine"),
2224 "FROM x AS x is an external pin, got: {:?}",
2225 c.nodes
2226 );
2227 }
2228
2229 const SAMPLE: &str = r"
2230use std::path::Path;
2231
2232pub struct Store;
2233
2234impl Store {
2235 pub fn open() -> Store {
2236 helper();
2237 Store
2238 }
2239}
2240
2241fn helper() {}
2242
2243mod inner {
2244 pub fn nested() {}
2245}
2246";
2247
2248 fn keys(fs: &crate::FactSet) -> Vec<String> {
2249 let mut k: Vec<_> = fs.nodes.iter().map(|n| n.key.clone()).collect();
2250 k.sort();
2251 k
2252 }
2253
2254 #[test]
2255 fn rust_extractor_emits_symbols_and_edges() {
2256 let fs = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
2257 let ks = keys(&fs);
2258 assert!(ks.contains(&"file:src/lib.rs".to_owned()));
2259 assert!(ks.contains(&"sym:rust:src/lib.rs#Store".to_owned()));
2260 assert!(ks.contains(&"sym:rust:src/lib.rs#Store::open".to_owned()));
2261 assert!(ks.contains(&"sym:rust:src/lib.rs#helper".to_owned()));
2262 assert!(ks.contains(&"sym:rust:src/lib.rs#inner".to_owned()));
2263 assert!(ks.contains(&"sym:rust:src/lib.rs#inner::nested".to_owned()));
2264
2265 let open = fs
2267 .nodes
2268 .iter()
2269 .find(|n| n.key == "sym:rust:src/lib.rs#Store::open")
2270 .expect("open node");
2271 assert_eq!(open.meta["calls"], serde_json::json!(["helper"]));
2272
2273 let defines: Vec<_> = fs
2275 .edges
2276 .iter()
2277 .filter(|e| e.kind == EdgeKind::Defines && e.dst == "sym:rust:src/lib.rs#helper")
2278 .collect();
2279 assert_eq!(defines.len(), 1);
2280 assert!(fs.edges.iter().any(|e| e.kind == EdgeKind::Contains
2281 && e.src == "sym:rust:src/lib.rs#inner"
2282 && e.dst == "sym:rust:src/lib.rs#inner::nested"));
2283
2284 assert!(fs.edges.iter().any(|e| e.kind == EdgeKind::Imports
2286 && e.src == "file:src/lib.rs"
2287 && e.dst == "import:rust:std::path::Path"));
2288 }
2289
2290 #[test]
2291 fn rust_extractor_records_struct_field_names() {
2292 let src = "pub struct ServeConfig {\n\
2295 \x20 pub addr: Option<String>,\n\
2296 \x20 pub tls_cert: Option<String>,\n\
2297 }\n\
2298 pub struct Pair(u8, u8);\n\
2299 pub struct Marker;\n";
2300 let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
2301 let fields = |key: &str| {
2302 fs.nodes
2303 .iter()
2304 .find(|n| n.key == key)
2305 .and_then(|n| n.meta.get("fields").cloned())
2306 };
2307 assert_eq!(
2308 fields("sym:rust:src/config.rs#ServeConfig"),
2309 Some(serde_json::json!(["addr", "tls_cert"])),
2310 "named fields captured in source order"
2311 );
2312 assert_eq!(fields("sym:rust:src/config.rs#Pair"), None);
2314 assert_eq!(fields("sym:rust:src/config.rs#Marker"), None);
2315 }
2316
2317 #[test]
2318 fn struct_records_field_types_and_config_root_marker() {
2319 let src = "// @rto:config\n\
2322 pub struct Config {\n\
2323 \x20 pub zerobus: ZerobusConfig,\n\
2324 \x20 pub replicas: Option<u32>,\n\
2325 }\n\
2326 pub struct ZerobusConfig {\n\
2327 \x20 pub server_endpoint: String,\n\
2328 }\n";
2329 let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
2330 let node = |key: &str| fs.nodes.iter().find(|n| n.key == key).expect("node");
2331 let root = node("sym:rust:src/config.rs#Config");
2332 assert_eq!(root.meta.get("config_root"), Some(&serde_json::json!(true)));
2333 assert_eq!(
2334 root.meta.get("field_types"),
2335 Some(&serde_json::json!({ "zerobus": "ZerobusConfig", "replicas": "u32" })),
2336 "transparent wrappers peeled (Option<u32> → u32)"
2337 );
2338 assert_eq!(
2340 node("sym:rust:src/config.rs#ZerobusConfig")
2341 .meta
2342 .get("config_root"),
2343 None
2344 );
2345 }
2346
2347 #[test]
2348 fn config_root_struct_synthesizes_recursive_dotted_config_keys() {
2349 let src = "// @rto:config\n\
2353 pub struct Config {\n\
2354 \x20 pub zerobus: ZerobusConfig,\n\
2355 \x20 pub log_level: String,\n\
2356 }\n\
2357 pub struct ZerobusConfig {\n\
2358 \x20 pub server_endpoint: String,\n\
2359 \x20 pub workspace_url: String,\n\
2360 }\n";
2361 let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
2362 let cfg = |dotted: &str| {
2363 fs.nodes
2364 .iter()
2365 .find(|n| n.key == format!("cfgkey:src/config.rs#{dotted}"))
2366 };
2367 for dotted in [
2368 "zerobus.server_endpoint",
2369 "zerobus.workspace_url",
2370 "log_level",
2371 ] {
2372 let n = cfg(dotted).unwrap_or_else(|| panic!("missing {dotted}: {:?}", fs.nodes));
2373 assert_eq!(n.kind, NodeKind::Other("config_key".into()));
2374 assert_eq!(n.meta.get("key").and_then(|v| v.as_str()), Some(dotted));
2375 assert_eq!(
2377 n.meta.get("source").and_then(|v| v.as_str()),
2378 Some("struct")
2379 );
2380 assert_eq!(
2381 n.meta.get("struct").and_then(|v| v.as_str()),
2382 Some("Config")
2383 );
2384 }
2385 assert!(
2387 cfg("zerobus").is_none(),
2388 "intermediate section is not a leaf"
2389 );
2390 assert!(fs.edges.iter().any(|e| e.src == "file:src/config.rs"
2392 && e.dst == "cfgkey:src/config.rs#zerobus.server_endpoint"
2393 && e.kind == EdgeKind::Contains));
2394 }
2395
2396 #[test]
2397 fn struct_without_config_marker_synthesizes_no_config_keys() {
2398 let src = "pub struct Config {\n\
2401 \x20 pub zerobus: ZerobusConfig,\n\
2402 }\n\
2403 pub struct ZerobusConfig {\n\
2404 \x20 pub server_endpoint: String,\n\
2405 }\n";
2406 let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
2407 assert!(
2408 fs.nodes
2409 .iter()
2410 .all(|n| n.kind != NodeKind::Other("config_key".into())),
2411 "no synthetic config_key nodes without the marker: {:?}",
2412 fs.nodes
2413 );
2414 }
2415
2416 #[test]
2417 fn config_root_recursion_terminates_on_a_type_cycle() {
2418 let src = "// @rto:config\n\
2421 pub struct Config {\n\
2422 \x20 pub addr: String,\n\
2423 \x20 pub next: Box<Config>,\n\
2424 }\n";
2425 let fs = RustExtractor.extract("src/config.rs", "b", src.as_bytes());
2426 let has = |dotted: &str| {
2427 fs.nodes
2428 .iter()
2429 .any(|n| n.key == format!("cfgkey:src/config.rs#{dotted}"))
2430 };
2431 assert!(has("addr"));
2432 assert!(has("next"), "cyclic field falls back to a leaf");
2435 assert!(!has("next.addr"), "no unbounded expansion");
2436 }
2437
2438 #[test]
2439 fn rust_extraction_is_deterministic() {
2440 let a = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
2441 let b = RustExtractor.extract("src/lib.rs", "blob1", SAMPLE.as_bytes());
2442 assert_eq!(a, b);
2443 }
2444
2445 #[test]
2446 fn rust_extractor_captures_doc_comments() {
2447 let src = "/// The central store.\n\
2448 pub struct Store;\n\n\
2449 /// Opens it.\n\
2450 /// Reads the config.\n\
2451 pub fn open() {}\n\n\
2452 // not a doc comment\n\
2453 pub fn plain() {}\n";
2454 let fs = RustExtractor.extract("src/lib.rs", "b", src.as_bytes());
2455 let content = |key: &str| {
2456 fs.nodes
2457 .iter()
2458 .find(|n| n.key == key)
2459 .and_then(|n| n.meta.get("content"))
2460 .and_then(|v| v.as_str())
2461 .map(ToOwned::to_owned)
2462 };
2463 assert_eq!(
2464 content("sym:rust:src/lib.rs#Store").as_deref(),
2465 Some("The central store.")
2466 );
2467 assert_eq!(
2468 content("sym:rust:src/lib.rs#open").as_deref(),
2469 Some("Opens it. Reads the config.")
2470 );
2471 assert_eq!(content("sym:rust:src/lib.rs#plain"), None);
2473 }
2474
2475 #[test]
2476 fn prose_file_captures_capped_body() {
2477 let md = FileNodeExtractor.extract("docs/x.md", "b", b"# Title\n\nSome prose here.\n");
2478 assert_eq!(md.nodes[0].meta["content"], "# Title Some prose here.");
2479 let rs = FileNodeExtractor.extract("notes.bin", "b", b"\x00\x01binary");
2481 assert!(rs.nodes[0].meta.get("content").is_none());
2482 let upper = FileNodeExtractor.extract("README.MD", "b", b"# Hi\n");
2484 assert_eq!(upper.nodes[0].meta["content"], "# Hi");
2485 }
2486
2487 #[cfg(feature = "pdf-text")]
2490 fn minimal_pdf(text: &str) -> Vec<u8> {
2491 let content = format!("BT /F1 24 Tf 72 720 Td ({text}) Tj ET");
2492 let objects = [
2493 "<< /Type /Catalog /Pages 2 0 R >>".to_owned(),
2494 "<< /Type /Pages /Kids [3 0 R] /Count 1 >>".to_owned(),
2495 "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>".to_owned(),
2496 format!("<< /Length {} >>\nstream\n{content}\nendstream", content.len()),
2497 "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>".to_owned(),
2498 ];
2499 let mut pdf = Vec::new();
2500 pdf.extend_from_slice(b"%PDF-1.4\n");
2501 let mut offsets = Vec::new();
2502 for (i, obj) in objects.iter().enumerate() {
2503 offsets.push(pdf.len());
2504 pdf.extend_from_slice(format!("{} 0 obj\n{obj}\nendobj\n", i + 1).as_bytes());
2505 }
2506 let xref_start = pdf.len();
2507 pdf.extend_from_slice(
2508 format!("xref\n0 {}\n0000000000 65535 f \n", objects.len() + 1).as_bytes(),
2509 );
2510 for off in &offsets {
2511 pdf.extend_from_slice(format!("{off:010} 00000 n \n").as_bytes());
2512 }
2513 pdf.extend_from_slice(
2514 format!(
2515 "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref_start}\n%%EOF\n",
2516 objects.len() + 1
2517 )
2518 .as_bytes(),
2519 );
2520 pdf
2521 }
2522
2523 #[cfg(feature = "pdf-text")]
2524 #[test]
2525 fn pdf_file_captures_text_content() {
2526 let pdf = minimal_pdf("Hello Roteiro");
2527 let facts = FileNodeExtractor.extract("docs/guide.pdf", "b", &pdf);
2528 let content = facts.nodes[0].meta["content"].as_str().unwrap();
2529 assert!(content.contains("Hello Roteiro"), "got: {content:?}");
2530 let upper = FileNodeExtractor.extract("docs/Guide.PDF", "b", &pdf);
2532 assert!(upper.nodes[0].meta.get("content").is_some());
2533 let bad = FileNodeExtractor.extract("docs/bad.pdf", "b", b"%PDF-1.4\ngarbage");
2535 assert!(bad.nodes[0].meta.get("content").is_none());
2536 }
2537
2538 #[cfg(any(feature = "image-ocr", feature = "image-vision"))]
2539 #[test]
2540 fn image_content_guards_before_touching_models() {
2541 assert!(super::is_image("shot.PNG"));
2543 assert!(super::is_image("b.jpeg"));
2544 assert!(super::is_image("c.jpg"));
2545 assert!(!super::is_image("d.gif"));
2546 assert!(
2548 super::image_content("notes.txt", b"hello", super::IngestConfig::default()).is_none()
2549 );
2550 let big = vec![0u8; super::MAX_IMAGE_BYTES + 1];
2552 assert!(super::image_content("shot.png", &big, super::IngestConfig::default()).is_none());
2553 }
2554
2555 #[test]
2556 fn doc_comment_body_recognises_doc_markers() {
2557 assert_eq!(super::doc_comment_body("/// hi").as_deref(), Some("hi"));
2558 assert_eq!(
2559 super::doc_comment_body("//! mod doc").as_deref(),
2560 Some("mod doc")
2561 );
2562 assert_eq!(
2563 super::doc_comment_body("/** block */").as_deref(),
2564 Some("block")
2565 );
2566 assert_eq!(super::doc_comment_body("// plain"), None);
2568 assert_eq!(super::doc_comment_body("//// header"), None);
2569 assert_eq!(super::doc_comment_body("/**/").as_deref(), Some(""));
2571 assert_eq!(super::doc_comment_body("/*!*/").as_deref(), Some(""));
2572 }
2573
2574 #[test]
2575 fn registry_dispatches_by_extension() {
2576 let rs = Registry::default().extract("src/lib.rs", "b", SAMPLE.as_bytes());
2577 assert!(rs.nodes.len() > 1, "rust file yields symbols");
2578 let txt = Registry::default().extract("notes.txt", "b", b"hello\n");
2579 assert_eq!(
2580 txt.nodes.len(),
2581 1,
2582 "non-code file falls back to a file node"
2583 );
2584 assert_eq!(txt.nodes[0].kind, NodeKind::File);
2585 }
2586
2587 #[test]
2588 fn tags_extracts_python_symbols_calls_and_nesting() {
2589 let src = "def helper():\n pass\n\nclass Thing:\n def run(self):\n helper()\n";
2590 let fs = Registry::default().extract("app.py", "b", src.as_bytes());
2591
2592 let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
2593 assert!(names.contains(&"helper"), "top-level function");
2594 assert!(names.contains(&"Thing"), "class");
2595 assert!(names.contains(&"run"), "method");
2596
2597 assert_eq!(
2599 fs.nodes
2600 .iter()
2601 .find(|n| n.name == "helper")
2602 .and_then(|n| n.lang.as_deref()),
2603 Some("python")
2604 );
2605
2606 assert!(
2608 fs.edges
2609 .iter()
2610 .any(|e| e.kind == EdgeKind::Contains && e.dst.ends_with("#Thing::run")),
2611 "method nested under class via containment"
2612 );
2613
2614 let run = fs.nodes.iter().find(|n| n.name == "run").unwrap();
2616 let calls = run.meta.get("calls").and_then(|v| v.as_array()).unwrap();
2617 assert!(
2618 calls.iter().any(|c| c.as_str() == Some("helper")),
2619 "enclosed call captured in meta.calls"
2620 );
2621 }
2622
2623 #[test]
2624 fn tags_extraction_is_deterministic() {
2625 let src = b"package main\nfunc Add(a int) int { return a }\n";
2626 let a = Registry::default().extract("m.go", "b", src);
2627 let b = Registry::default().extract("m.go", "b", src);
2628 assert_eq!(a, b, "tags extraction must be deterministic");
2629 assert!(
2630 a.nodes
2631 .iter()
2632 .any(|n| n.name == "Add" && n.kind == NodeKind::Fn)
2633 );
2634 }
2635
2636 #[test]
2637 fn tags_extracts_typescript() {
2638 let ts = Registry::default().extract("svc.ts", "b", b"export class Svc {\n run() {}\n}\n");
2639 assert!(ts.nodes.iter().any(|n| n.name == "Svc"), "class");
2640 assert!(ts.nodes.iter().any(|n| n.name == "run"), "method");
2641 assert_eq!(
2642 ts.nodes
2643 .iter()
2644 .find(|n| n.name == "Svc")
2645 .and_then(|n| n.lang.as_deref()),
2646 Some("typescript")
2647 );
2648 }
2649
2650 fn import_targets(path: &str, src: &[u8]) -> Vec<String> {
2654 Registry::default()
2655 .extract(path, "b", src)
2656 .nodes
2657 .iter()
2658 .filter(|n| n.kind == NodeKind::Other("import".into()))
2659 .inspect(|n| {
2660 assert!(
2661 n.path.is_none(),
2662 "import node must not be file-scoped: {}",
2663 n.key
2664 );
2665 })
2666 .map(|n| n.key.clone())
2667 .collect()
2668 }
2669
2670 #[test]
2671 fn extracts_imports_edges_per_language() {
2672 let cases: &[(&str, &[u8], &[&str])] = &[
2675 (
2676 "app.py",
2677 b"import os\nfrom a.b import c\nimport x.y as z\n",
2678 &["import:python:os", "import:python:a.b", "import:python:x.y"],
2679 ),
2680 (
2681 "m.js",
2682 b"import foo from \"./mod.js\";\nexport { y } from \"./y.js\";\n",
2683 &["import:javascript:./mod.js", "import:javascript:./y.js"],
2684 ),
2685 (
2686 "svc.ts",
2687 b"import { A } from \"./a\";\n",
2688 &["import:typescript:./a"],
2689 ),
2690 (
2691 "m.go",
2692 b"package main\nimport (\n\t\"fmt\"\n\t\"os\"\n)\n",
2693 &["import:go:fmt", "import:go:os"],
2694 ),
2695 (
2696 "M.java",
2697 b"import java.util.List;\nimport static a.B.c;\n",
2698 &["import:java:java.util.List", "import:java:a.B.c"],
2699 ),
2700 (
2701 "m.c",
2702 b"#include <stdio.h>\n#include \"local.h\"\n",
2703 &["import:c:stdio.h", "import:c:local.h"],
2704 ),
2705 ("m.cpp", b"#include <vector>\n", &["import:cpp:vector"]),
2706 ];
2707 for (path, src, expected) in cases {
2708 let got = import_targets(path, src);
2709 for want in *expected {
2710 assert!(
2711 got.iter().any(|k| k == want),
2712 "{path}: expected import node {want}, got {got:?}"
2713 );
2714 }
2715 let fs = Registry::default().extract(path, "b", src);
2717 for want in *expected {
2718 assert!(
2719 fs.edges.iter().any(|e| e.kind == EdgeKind::Imports
2720 && e.src == format!("file:{path}")
2721 && &e.dst == want),
2722 "{path}: expected Imports edge to {want}"
2723 );
2724 }
2725 }
2726 }
2727
2728 #[test]
2729 fn every_registered_language_query_compiles() {
2730 for ext in [
2734 "py", "js", "ts", "tsx", "go", "rb", "java", "c", "cpp", "cs", "php", "scala", "ml",
2735 "mli", "ex", "sh", "sql",
2736 ] {
2737 let def = super::tag_lang_for(ext).unwrap_or_else(|| panic!("no language for .{ext}"));
2738 let lang = def.lang;
2739 assert!(
2740 super::tag_config(&def).is_some(),
2741 "tags query for .{ext} ({lang}) must compile against its grammar"
2742 );
2743 }
2744 }
2745
2746 #[test]
2747 fn ocaml_impl_and_interface_cache_under_distinct_grammars() {
2748 let ml = super::tag_lang_for("ml").unwrap();
2752 let mli = super::tag_lang_for("mli").unwrap();
2753 assert_eq!(ml.lang, "ocaml");
2754 assert_eq!(mli.lang, "ocaml");
2755 assert_ne!(
2756 ml.grammar_key, mli.grammar_key,
2757 "distinct grammars must cache separately"
2758 );
2759 }
2760
2761 #[test]
2762 fn tags_extracts_vendored_bash_query() {
2763 let src = "greet() {\n echo hi\n}\nmain() {\n greet\n}\n";
2764 let fs = Registry::default().extract("run.sh", "b", src.as_bytes());
2765 let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
2766 assert!(names.contains(&"greet"), "shell function greet");
2767 assert!(names.contains(&"main"), "shell function main");
2768
2769 let main = fs.nodes.iter().find(|n| n.name == "main").unwrap();
2771 assert!(
2772 main.meta
2773 .get("calls")
2774 .and_then(|v| v.as_array())
2775 .is_some_and(|c| c.iter().any(|x| x.as_str() == Some("greet"))),
2776 "internal command invocation captured"
2777 );
2778 }
2779
2780 #[test]
2781 fn tags_extracts_vendored_sql_query() {
2782 let src = "CREATE TABLE users (id int);\n\
2783 CREATE FUNCTION recent() RETURNS int AS $$ SELECT total(id) FROM users $$ LANGUAGE sql;\n";
2784 let fs = Registry::default().extract("schema.sql", "b", src.as_bytes());
2785 let names: Vec<&str> = fs.nodes.iter().map(|n| n.name.as_str()).collect();
2786 assert!(names.contains(&"users"), "table definition");
2787 assert!(names.contains(&"recent"), "function definition");
2788
2789 assert_eq!(
2791 fs.nodes.iter().find(|n| n.name == "users").map(|n| &n.kind),
2792 Some(&NodeKind::Other("table".to_owned()))
2793 );
2794 let f = fs.nodes.iter().find(|n| n.name == "recent").unwrap();
2796 assert!(
2797 f.meta
2798 .get("calls")
2799 .and_then(|v| v.as_array())
2800 .is_some_and(|c| c.iter().any(|x| x.as_str() == Some("total"))),
2801 "invocation inside function captured in meta.calls"
2802 );
2803 assert_eq!(
2804 fs.nodes
2805 .iter()
2806 .find(|n| n.name == "users")
2807 .and_then(|n| n.lang.as_deref()),
2808 Some("sql")
2809 );
2810 }
2811
2812 #[test]
2813 fn ingest_prose_toggle_gates_embedded_content() {
2814 use super::IngestConfig;
2815
2816 let content = |ingest: IngestConfig| {
2817 Registry::new(ingest)
2818 .extract("notes.md", "b", b"# Title\n\nBody text.\n")
2819 .nodes[0]
2820 .meta
2821 .get("content")
2822 .and_then(|v| v.as_str())
2823 .map(str::to_owned)
2824 };
2825
2826 assert!(
2828 content(IngestConfig::default()).is_some_and(|c| c.contains("Body text")),
2829 "prose content embedded by default"
2830 );
2831 assert_eq!(
2832 content(IngestConfig {
2833 prose: false,
2834 ..IngestConfig::default()
2835 }),
2836 None,
2837 "disabling prose suppresses the embedded body"
2838 );
2839 }
2840
2841 #[test]
2842 fn env_tag_stable_by_default_and_shifts_when_gated() {
2843 use super::IngestConfig;
2844
2845 let all_on = Registry::new(IngestConfig::default()).env_tag();
2848 assert_eq!(all_on, Registry::default().env_tag());
2849
2850 let no_prose = Registry::new(IngestConfig {
2853 prose: false,
2854 ..IngestConfig::default()
2855 })
2856 .env_tag();
2857 let no_pdf = Registry::new(IngestConfig {
2858 pdf: false,
2859 ..IngestConfig::default()
2860 })
2861 .env_tag();
2862 let no_audio = Registry::new(IngestConfig {
2863 audio: false,
2864 ..IngestConfig::default()
2865 })
2866 .env_tag();
2867 assert_ne!(no_prose, all_on);
2868 assert_ne!(no_pdf, all_on);
2869 assert_ne!(no_audio, all_on);
2870 assert_ne!(no_prose, no_pdf);
2871 assert_ne!(no_audio, no_prose);
2872 assert_ne!(no_audio, no_pdf);
2873 }
2874}