From a235911a12608e9aeaf90ee89fa8430841da8c37 Mon Sep 17 00:00:00 2001
From: forhappy <haipingf@gmail.com>
Date: Thu, 23 Jul 2026 19:12:04 -0700
Subject: [PATCH] Harden R language extraction parity

---
 crates/compass-languages/src/engine.rs     |   1 +
 crates/compass-languages/src/lib.rs        |   1 +
 crates/compass-languages/src/r.rs          | 363 +++++++++++++++++++++
 crates/compass-languages/src/registry.rs   |  12 +-
 crates/compass-languages/tests/registry.rs |   4 +-
 crates/compass-parity/src/lib.rs           |  56 +++-
 6 files changed, 433 insertions(+), 4 deletions(-)
 create mode 100644 crates/compass-languages/src/r.rs

diff --git a/crates/compass-languages/src/engine.rs b/crates/compass-languages/src/engine.rs
index fe7d434..0edd938 100644
--- a/crates/compass-languages/src/engine.rs
+++ b/crates/compass-languages/src/engine.rs
@@ -122,6 +122,7 @@ impl Engine {
             "zig" => return Ok(crate::zig::extract(path, source)),
             "verilog" => return Ok(crate::verilog::extract(path, source)),
             "sql" => return Ok(crate::sql::extract(path, source)),
+            "r" => return Ok(crate::r::extract(path, source)),
             "pascal" => return Ok(crate::pascal::extract(path, source)),
             "apex" => return Ok(crate::apex::extract(path, source)),
             "dart" => return Ok(crate::dart::extract(path, source)),
diff --git a/crates/compass-languages/src/lib.rs b/crates/compass-languages/src/lib.rs
index bc22044..af5cf09 100644
--- a/crates/compass-languages/src/lib.rs
+++ b/crates/compass-languages/src/lib.rs
@@ -27,6 +27,7 @@ mod pascal_forms;
 mod php;
 mod powershell;
 mod program;
+mod r;
 mod registry;
 mod rust_lang;
 mod scip;
diff --git a/crates/compass-languages/src/r.rs b/crates/compass-languages/src/r.rs
new file mode 100644
index 0000000..b757a45
--- /dev/null
+++ b/crates/compass-languages/src/r.rs
@@ -0,0 +1,363 @@
+use std::collections::{HashMap, HashSet};
+use std::path::Path;
+
+use compass_model::{EdgeRecord, NodeRecord};
+use regex::Regex;
+use serde_json::{Map, Value};
+
+use crate::{Extraction, RawCall, file_stem, make_id};
+
+const NON_CALLS: &[&str] = &[
+    "function", "if", "for", "while", "repeat", "switch", "return",
+];
+
+pub(crate) fn extract(path: &Path, source: &[u8]) -> Extraction {
+    State::new(path, source).run()
+}
+
+#[derive(Clone)]
+struct Function {
+    start: usize,
+    end: usize,
+    id: String,
+    name: String,
+    parent: Option<String>,
+}
+
+struct State<'a> {
+    path: &'a Path,
+    text: String,
+    masked: Vec<u8>,
+    source_file: String,
+    stem: String,
+    file_id: String,
+    extraction: Extraction,
+    seen: HashSet<String>,
+}
+
+impl<'a> State<'a> {
+    fn new(path: &'a Path, source: &'a [u8]) -> Self {
+        let source_file = path.to_string_lossy().into_owned();
+        let text = String::from_utf8_lossy(source).into_owned();
+        let masked = mask_non_code(text.as_bytes());
+        Self {
+            path,
+            text,
+            masked,
+            stem: file_stem(path),
+            file_id: make_id(&[&source_file]),
+            source_file,
+            extraction: Extraction::default(),
+            seen: HashSet::new(),
+        }
+    }
+
+    fn run(mut self) -> Extraction {
+        let label = self
+            .path
+            .file_name()
+            .and_then(|name| name.to_str())
+            .unwrap_or_default();
+        self.add_node(&self.file_id.clone(), label, 1);
+        self.add_imports();
+
+        let functions = self.functions();
+        for function in &functions {
+            let at = self.line_at(function.start);
+            self.add_node(&function.id, &format!("{}()", function.name), at);
+            let container = function
+                .parent
+                .as_deref()
+                .unwrap_or(&self.file_id)
+                .to_owned();
+            self.add_edge(&container, &function.id, "contains", at);
+        }
+
+        let mut labels: HashMap<String, Vec<Function>> = HashMap::new();
+        for function in &functions {
+            labels
+                .entry(function.name.clone())
+                .or_default()
+                .push(function.clone());
+        }
+        let mut seen_calls = HashSet::new();
+        for function in &functions {
+            self.add_calls(function, &functions, &labels, &mut seen_calls);
+        }
+
+        self.extraction.edges.retain(|edge| {
+            self.seen.contains(&edge.source)
+                && (self.seen.contains(&edge.target)
+                    || edge.attributes.get("relation").and_then(Value::as_str) == Some("imports"))
+        });
+        self.extraction
+    }
+
+    fn add_imports(&mut self) {
+        let Ok(imports) = Regex::new(r"(?m)^[ \t]*(?:library|require|requireNamespace)[ \t]*\(")
+        else {
+            return;
+        };
+        let Ok(package) = Regex::new(r#"^\s*(?:package\s*=\s*)?["']?([A-Za-z][A-Za-z0-9._-]*)"#)
+        else {
+            return;
+        };
+        let masked = String::from_utf8_lossy(&self.masked);
+        let mut found = Vec::new();
+        for matched in imports.find_iter(&masked) {
+            let Some(arguments_end) = matching_delimiter(&self.masked, matched.end(), b'(', b')')
+            else {
+                continue;
+            };
+            if let Some(capture) = package.captures(&self.text[matched.end()..arguments_end])
+                && let Some(name) = capture.get(1)
+            {
+                found.push((matched.start(), name.as_str().to_owned()));
+            }
+        }
+        for (offset, name) in found {
+            self.add_edge(
+                &self.file_id.clone(),
+                &make_id(&[&name]),
+                "imports",
+                self.line_at(offset),
+            );
+        }
+    }
+
+    fn functions(&self) -> Vec<Function> {
+        let Ok(pattern) = Regex::new(
+            r"(?m)^[ \t]*([A-Za-z.][A-Za-z0-9._]*)[ \t]*(?:<<-|<-|=)[ \t]*function[ \t]*\(",
+        ) else {
+            return Vec::new();
+        };
+        let masked = String::from_utf8_lossy(&self.masked);
+        let specs = pattern
+            .captures_iter(&masked)
+            .filter_map(|capture| {
+                let whole = capture.get(0)?;
+                let name = capture.get(1)?.as_str().to_owned();
+                let parameters_end = matching_delimiter(&self.masked, whole.end(), b'(', b')')?;
+                let mut body_start = parameters_end + 1;
+                while self
+                    .masked
+                    .get(body_start)
+                    .is_some_and(|byte| byte.is_ascii_whitespace())
+                {
+                    body_start += 1;
+                }
+                let end = if self.masked.get(body_start) == Some(&b'{') {
+                    matching_delimiter(&self.masked, body_start + 1, b'{', b'}')?
+                } else {
+                    self.masked[body_start..]
+                        .iter()
+                        .position(|byte| *byte == b'\n')
+                        .map_or_else(
+                            || self.masked.len().saturating_sub(1),
+                            |relative| body_start + relative,
+                        )
+                };
+                Some((whole.start(), end, name))
+            })
+            .collect::<Vec<_>>();
+        let mut functions: Vec<Function> = Vec::new();
+        for (start, end, name) in specs {
+            let parent = functions
+                .iter()
+                .filter(|function| function.start < start && end <= function.end)
+                .min_by_key(|function| function.end - function.start)
+                .map(|function| function.id.clone());
+            let id = make_id(&[parent.as_deref().unwrap_or(&self.stem), &name]);
+            functions.push(Function {
+                start,
+                end,
+                id,
+                name,
+                parent,
+            });
+        }
+        functions
+    }
+
+    fn add_calls(
+        &mut self,
+        function: &Function,
+        functions: &[Function],
+        labels: &HashMap<String, Vec<Function>>,
+        seen: &mut HashSet<(String, String)>,
+    ) {
+        let Ok(calls) =
+            Regex::new(r"([A-Za-z.][A-Za-z0-9._]*(?:(?:::|\$)[A-Za-z.][A-Za-z0-9._]*)?)[ \t]*\(")
+        else {
+            return;
+        };
+        let mut body = self.masked[function.start..=function.end].to_vec();
+        for nested in functions {
+            if function.start < nested.start && nested.end <= function.end {
+                for index in nested.start..=nested.end {
+                    if body[index - function.start] != b'\n' {
+                        body[index - function.start] = b' ';
+                    }
+                }
+            }
+        }
+        let body = String::from_utf8_lossy(&body);
+        for capture in calls.captures_iter(&body) {
+            let Some(raw) = capture.get(1) else {
+                continue;
+            };
+            let expression = raw.as_str();
+            let callee = expression
+                .rsplit_once("::")
+                .or_else(|| expression.rsplit_once('$'))
+                .map_or(expression, |(_, name)| name);
+            if NON_CALLS.contains(&callee) {
+                continue;
+            }
+            let absolute = function.start + raw.start();
+            if let Some(target) = call_target(callee, function, labels) {
+                if seen.insert((function.id.clone(), target.clone())) {
+                    self.add_edge(&function.id, &target, "calls", self.line_at(absolute));
+                }
+            } else {
+                let source_location = format!("L{}", self.line_at(absolute));
+                self.extraction.raw_calls_mut().push(RawCall {
+                    caller_nid: function.id.clone(),
+                    callee: callee.to_owned(),
+                    is_member_call: Some(expression.contains('$')),
+                    source_file: self.source_file.clone(),
+                    source_location,
+                    receiver: None,
+                    receiver_type: None,
+                    lang: None,
+                    extensions: Map::new(),
+                });
+            }
+        }
+    }
+
+    fn add_node(&mut self, id: &str, label: &str, at: usize) {
+        if !self.seen.insert(id.to_owned()) {
+            return;
+        }
+        let mut attributes = Map::new();
+        attributes.insert("label".into(), Value::String(label.to_owned()));
+        attributes.insert("file_type".into(), Value::String("code".into()));
+        attributes.insert(
+            "source_file".into(),
+            Value::String(self.source_file.clone()),
+        );
+        attributes.insert("source_location".into(), Value::String(format!("L{at}")));
+        self.extraction.nodes.push(NodeRecord {
+            id: id.to_owned(),
+            attributes,
+        });
+    }
+
+    fn add_edge(&mut self, source: &str, target: &str, relation: &str, at: usize) {
+        let mut attributes = Map::new();
+        attributes.insert("relation".into(), Value::String(relation.to_owned()));
+        attributes.insert("confidence".into(), Value::String("EXTRACTED".into()));
+        attributes.insert(
+            "source_file".into(),
+            Value::String(self.source_file.clone()),
+        );
+        attributes.insert("source_location".into(), Value::String(format!("L{at}")));
+        attributes.insert("weight".into(), Value::from(1.0));
+        self.extraction.edges.push(EdgeRecord {
+            source: source.to_owned(),
+            target: target.to_owned(),
+            attributes,
+        });
+    }
+
+    fn line_at(&self, offset: usize) -> usize {
+        self.text.as_bytes()[..offset.min(self.text.len())]
+            .iter()
+            .filter(|byte| **byte == b'\n')
+            .count()
+            + 1
+    }
+}
+
+fn call_target(
+    callee: &str,
+    caller: &Function,
+    labels: &HashMap<String, Vec<Function>>,
+) -> Option<String> {
+    let candidates = labels
+        .get(callee)?
+        .iter()
+        .filter(|candidate| candidate.id != caller.id)
+        .collect::<Vec<_>>();
+    let children = candidates
+        .iter()
+        .filter(|candidate| candidate.parent.as_deref() == Some(caller.id.as_str()))
+        .collect::<Vec<_>>();
+    if children.len() == 1 {
+        return Some(children[0].id.clone());
+    }
+    let same_scope = candidates
+        .iter()
+        .filter(|candidate| candidate.parent == caller.parent)
+        .collect::<Vec<_>>();
+    if same_scope.len() == 1 {
+        return Some(same_scope[0].id.clone());
+    }
+    (candidates.len() == 1).then(|| candidates[0].id.clone())
+}
+
+fn mask_non_code(source: &[u8]) -> Vec<u8> {
+    let mut masked = source.to_vec();
+    let mut quote = None;
+    let mut escaped = false;
+    let mut comment = false;
+    for (index, byte) in source.iter().copied().enumerate() {
+        if comment {
+            if byte == b'\n' {
+                comment = false;
+            } else {
+                masked[index] = b' ';
+            }
+            continue;
+        }
+        if let Some(active) = quote {
+            if byte == b'\n' {
+                continue;
+            } else if escaped {
+                escaped = false;
+                masked[index] = b' ';
+            } else if byte == b'\\' {
+                escaped = true;
+                masked[index] = b' ';
+            } else if byte == active {
+                quote = None;
+            } else {
+                masked[index] = b' ';
+            }
+            continue;
+        }
+        if byte == b'#' {
+            comment = true;
+            masked[index] = b' ';
+        } else if matches!(byte, b'\'' | b'"') {
+            quote = Some(byte);
+        }
+    }
+    masked
+}
+
+fn matching_delimiter(source: &[u8], start: usize, opening: u8, closing: u8) -> Option<usize> {
+    let mut depth = 1_u32;
+    for (index, byte) in source.iter().copied().enumerate().skip(start) {
+        if byte == opening {
+            depth += 1;
+        } else if byte == closing {
+            depth -= 1;
+            if depth == 0 {
+                return Some(index);
+            }
+        }
+    }
+    None
+}
diff --git a/crates/compass-languages/src/registry.rs b/crates/compass-languages/src/registry.rs
index 8a6df7a..de4db39 100644
--- a/crates/compass-languages/src/registry.rs
+++ b/crates/compass-languages/src/registry.rs
@@ -99,6 +99,11 @@ impl Registry {
             "dart" => spec("dart", "dart", ExtractorKind::Generic),
             "v" | "sv" | "svh" => spec("verilog", "verilog", ExtractorKind::Generic),
             "sql" => spec("sql", "sql", ExtractorKind::Generic),
+            "r" => LanguageSpec {
+                name: "r",
+                grammar: None,
+                kind: ExtractorKind::Generic,
+            },
             "md" | "mdx" | "qmd" | "skill" => LanguageSpec {
                 name: "markdown",
                 grammar: None,
@@ -152,7 +157,7 @@ impl Registry {
         &[
             "py", "ts", "tsx", "js", "go", "rs", "java", "c", "cpp", "rb", "cs", "kt", "scala",
             "php", "swift", "lua", "zig", "ps1", "ex", "m", "jl", "f90", "vue", "svelte", "astro",
-            "dart", "v", "sql", "md", "pas", "dfm", "sh", "json", "tf", "dm", "sln", "csproj",
+            "dart", "v", "sql", "r", "md", "pas", "dfm", "sh", "json", "tf", "dm", "sln", "csproj",
             "xaml", "razor", "cls",
         ]
     }
@@ -182,6 +187,11 @@ fn shebang_spec(path: &Path) -> Option<LanguageSpec> {
         "lua" => Some(spec("lua", "lua", ExtractorKind::Generic)),
         "php" => Some(spec("php", "php", ExtractorKind::Generic)),
         "julia" => Some(spec("julia", "julia", ExtractorKind::Generic)),
+        "Rscript" => Some(LanguageSpec {
+            name: "r",
+            grammar: None,
+            kind: ExtractorKind::Generic,
+        }),
         _ => None,
     }
 }
diff --git a/crates/compass-languages/tests/registry.rs b/crates/compass-languages/tests/registry.rs
index 81b9e10..b0c0528 100644
--- a/crates/compass-languages/tests/registry.rs
+++ b/crates/compass-languages/tests/registry.rs
@@ -12,8 +12,8 @@ fn registry_covers_every_python_dispatch_extension() -> Result<(), Box<dyn Error
         "gradle", "c", "cpp", "cc", "cxx", "hpp", "cu", "cuh", "metal", "rb", "rake", "cs", "kt",
         "kts", "scala", "php", "swift", "lua", "luau", "toc", "zig", "ps1", "psm1", "psd1", "ex",
         "exs", "mm", "jl", "f", "F", "f90", "F90", "f95", "F95", "f03", "F03", "f08", "F08", "vue",
-        "svelte", "astro", "dart", "v", "sv", "svh", "sql", "md", "mdx", "qmd", "skill", "pas",
-        "pp", "dpr", "dpk", "lpr", "inc", "dfm", "lfm", "lpk", "sh", "bash", "json", "tf",
+        "svelte", "astro", "dart", "v", "sv", "svh", "sql", "r", "md", "mdx", "qmd", "skill",
+        "pas", "pp", "dpr", "dpk", "lpr", "inc", "dfm", "lfm", "lpk", "sh", "bash", "json", "tf",
         "tfvars", "hcl", "dm", "dme", "dmi", "dmm", "dmf", "sln", "slnx", "csproj", "fsproj",
         "vbproj", "xaml", "razor", "cshtml", "cls", "trigger",
     ];
diff --git a/crates/compass-parity/src/lib.rs b/crates/compass-parity/src/lib.rs
index 1d30923..bd151f1 100644
--- a/crates/compass-parity/src/lib.rs
+++ b/crates/compass-parity/src/lib.rs
@@ -20,7 +20,7 @@ mod tests {
         deduplicate_entities, find_import_cycles, god_nodes, graph_diff, label_communities_by_hub,
         score_communities, suggest_questions, surprising_connections,
     };
-    use compass_languages::Engine;
+    use compass_languages::{Engine, Registry};
     use compass_mcp::GraphifyMcp;
     use compass_media::{docx_to_markdown, extract_pdf_text, xlsx_to_markdown};
     use compass_output::{
@@ -2228,6 +2228,40 @@ print(json.dumps({'content': content, 'default': default, 'omitted': omitted}, e
         compare_extraction("sample.py", "extract_python")
     }
 
+    #[test]
+    fn compass_registry_covers_graphify_dispatch_surface() -> Result<(), Box<dyn Error>> {
+        let repo = repository_root();
+        let output = Command::new(python_executable(&repo))
+            .args([
+                "-c",
+                "import json; from graphify.extract import _DISPATCH; print(json.dumps(sorted(_DISPATCH)))",
+            ])
+            .current_dir(&repo)
+            .env("PYTHONPATH", &repo)
+            .output()?;
+        assert!(
+            output.status.success(),
+            "{}",
+            String::from_utf8_lossy(&output.stderr)
+        );
+        let extensions: Vec<String> = serde_json::from_slice(&output.stdout)?;
+        let directory = tempfile::tempdir()?;
+        for extension in extensions {
+            let path = directory.path().join(format!("sample{extension}"));
+            let source = if extension.eq_ignore_ascii_case(".m") {
+                "@implementation Compass\n@end\n"
+            } else {
+                ""
+            };
+            fs::write(&path, source)?;
+            assert!(
+                Registry::resolve(&path).is_some(),
+                "Compass registry is missing Graphify dispatch extension {extension}"
+            );
+        }
+        Ok(())
+    }
+
     #[test]
     fn python_rationale_extraction_matches_exactly() -> Result<(), Box<dyn Error>> {
         let directory = tempfile::tempdir()?;
@@ -2568,6 +2602,19 @@ output "instance_id" { value = aws_instance.web.id }
         Ok(())
     }
 
+    #[test]
+    fn r_extraction_matches_exactly() -> Result<(), Box<dyn Error>> {
+        compare_extraction("sample.r", "extract_r")?;
+        let directory = tempfile::tempdir()?;
+        let invalid_utf8 = directory.path().join("legacy.r");
+        fs::write(
+            &invalid_utf8,
+            b"# legacy annotation: \xff\nrun <- function(value) value\n",
+        )?;
+        compare_extraction_path(&invalid_utf8, "extract_r")?;
+        Ok(())
+    }
+
     #[test]
     fn pascal_source_extraction_matches_exactly() -> Result<(), Box<dyn Error>> {
         compare_extraction("sample.pas", "extract_pascal")?;
@@ -2832,6 +2879,13 @@ hydrate();
             "#!/usr/bin/node\nfunction serve() { console.log('ok') }\nserve()\n",
         )?;
         compare_extraction_path(&node, "extract_js")?;
+
+        let rscript = directory.path().join("analyze");
+        fs::write(
+            &rscript,
+            "#!/usr/bin/env Rscript\nrun <- function(values) {\n  mean(values)\n}\n",
+        )?;
+        compare_extraction_path(&rscript, "extract_r")?;
         Ok(())
     }