improvise/crates/improvise-protocol/src/source_scan.rs

//! Production Rust source, for the guards that read source text: no comments,
//! no literals, no test code. The one home of that rule, `context/code/protocol.md` §2.

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};

/// Every `.rs` file under `roots` that is compiled outside tests, with its
/// comments, literals and test-only items blanked or dropped.
pub fn production_code(roots: &[PathBuf]) -> Vec<(PathBuf, String)> {
    let mut files = Vec::new();
    for root in roots {
        collect_rust(root, &mut files);
    }
    let code: Vec<(PathBuf, String)> = files
        .into_iter()
        .map(|file| {
            let text = std::fs::read_to_string(&file).expect("a listed source reads");
            let code = code_only(&text);
            (file, code)
        })
        .collect();
    let test_modules: BTreeSet<PathBuf> = code
        .iter()
        .flat_map(|(file, code)| test_module_paths(file, code))
        .collect();
    code.into_iter()
        .filter(|(file, _)| !test_modules.iter().any(|m| is_module_file(m, file)))
        .map(|(file, code)| (file, strip_test_items(&code)))
        .collect()
}

/// One file's text as [`production_code`] cleans it. Which *files* are test
/// modules needs their declaring file too, so only `production_code` knows.
pub fn production_text(text: &str) -> String {
    strip_test_items(&code_only(text))
}

fn collect_rust(dir: &Path, into: &mut Vec<PathBuf>) {
    let Ok(entries) = std::fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        match path.is_dir() {
            true => collect_rust(&path, into),
            false if path.extension().is_some_and(|e| e == "rs") => into.push(path),
            false => {}
        }
    }
}

/// `m` is a module path without extension: its file is `m.rs`, `m/mod.rs`,
/// or anything under `m/`.
fn is_module_file(m: &Path, file: &Path) -> bool {
    file == m.with_extension("rs") || file.starts_with(m)
}

/// The out-of-line modules `file` declares under a test-only `cfg`, as paths
/// without extension (`#[cfg(test)] mod x;` is `<dir>/x`).
fn test_module_paths(file: &Path, code: &str) -> Vec<PathBuf> {
    let dir = match file.file_stem().and_then(|s| s.to_str()) {
        Some("lib" | "main" | "mod") => file.parent().map(Path::to_path_buf),
        Some(stem) => file.parent().map(|p| p.join(stem)),
        None => None,
    }
    .expect("a source file has a directory and a stem");
    let lines: Vec<&str> = code.lines().map(str::trim).collect();
    let mut found = Vec::new();
    for (i, line) in lines.iter().enumerate() {
        if !is_test_cfg(line) {
            continue;
        }
        let item = lines[i + 1..]
            .iter()
            .find(|l| !l.starts_with("#[") && !l.is_empty());
        if let Some(name) = item.and_then(|l| out_of_line_mod(l)) {
            found.push(dir.join(name));
        }
    }
    found
}

/// `mod x;`, `pub mod x;`, `pub(crate) mod x;` → `x`.
fn out_of_line_mod(line: &str) -> Option<&str> {
    let rest = line.split_once("mod ")?.1.strip_suffix(';')?;
    let prefix = line.split_once("mod ")?.0.trim();
    (prefix.is_empty() || prefix.starts_with("pub")).then_some(rest.trim())
}

/// A `cfg` attribute that compiles its item only under `test`. `cfg(not(test))`
/// is production, and `"test-support"` is a literal, already blanked.
fn is_test_cfg(line: &str) -> bool {
    let line = line.trim();
    line.starts_with("#[cfg(")
        && !line.contains("not(")
        && line
            .split(|c: char| !c.is_alphanumeric() && c != '_')
            .any(|word| word == "test")
}

/// Drop every test-only item. rustfmt puts an item's closing brace at the
/// attribute's own indentation, which is what makes this a line scan (§2).
fn strip_test_items(text: &str) -> String {
    let lines: Vec<&str> = text.lines().collect();
    let mut kept = String::new();
    let mut i = 0;
    while i < lines.len() {
        let line = lines[i];
        match is_test_cfg(line) {
            false => {
                kept.push_str(line);
                kept.push('\n');
                i += 1;
            }
            true => {
                let indent = line.len() - line.trim_start().len();
                let close = format!("{}}}", " ".repeat(indent));
                i += 1;
                while i < lines.len() && lines[i].trim_start().starts_with("#[") {
                    i += 1;
                }
                match lines.get(i).is_some_and(|l| l.trim_end().ends_with(';')) {
                    true => i += 1,
                    false => {
                        while i < lines.len() && lines[i].trim_end() != close {
                            i += 1;
                        }
                        i += 1;
                    }
                }
            }
        }
    }
    kept
}

/// `text` with comments, string, byte-string, raw-string and char literals
/// replaced by spaces. Newlines survive, so line-based scans still line up.
pub fn code_only(text: &str) -> String {
    let chars: Vec<char> = text.chars().collect();
    let mut out = String::with_capacity(text.len());
    let mut i = 0;
    while i < chars.len() {
        let end = literal_or_comment_end(&chars, i);
        match end {
            Some(end) => {
                for &c in &chars[i..end] {
                    out.push(if c == '\n' { '\n' } else { ' ' });
                }
                i = end;
            }
            None => {
                out.push(chars[i]);
                i += 1;
            }
        }
    }
    out
}

/// Where the comment or literal starting at `i` ends, if one starts there.
fn literal_or_comment_end(c: &[char], i: usize) -> Option<usize> {
    let at = |k: usize| c.get(k).copied();
    let word_before = i > 0 && (c[i - 1].is_alphanumeric() || c[i - 1] == '_');
    match (at(i), at(i + 1)) {
        (Some('/'), Some('/')) => Some(line_end(c, i)),
        (Some('/'), Some('*')) => Some(block_end(c, i)),
        (Some('"'), _) => Some(quoted_end(c, i + 1, '"')),
        (Some('b'), Some('"')) if !word_before => Some(quoted_end(c, i + 2, '"')),
        (Some('r' | 'b'), Some('r' | '#' | '"')) if !word_before => raw_end(c, i),
        (Some('\''), _) => char_end(c, i),
        _ => None,
    }
}

fn line_end(c: &[char], from: usize) -> usize {
    (from..c.len()).find(|&k| c[k] == '\n').unwrap_or(c.len())
}

/// Block comments nest in Rust.
fn block_end(c: &[char], from: usize) -> usize {
    let mut depth = 0usize;
    let mut k = from;
    while k + 1 < c.len() {
        match (c[k], c[k + 1]) {
            ('/', '*') => {
                depth += 1;
                k += 2;
            }
            ('*', '/') => {
                depth -= 1;
                k += 2;
                if depth == 0 {
                    return k;
                }
            }
            _ => k += 1,
        }
    }
    c.len()
}

fn quoted_end(c: &[char], mut k: usize, quote: char) -> usize {
    while k < c.len() {
        match c[k] {
            '\\' => k += 2,
            q if q == quote => return k + 1,
            _ => k += 1,
        }
    }
    c.len()
}

/// `r"…"`, `r#"…"#`, `br#"…"#`. `None` when the `r` is only an identifier.
fn raw_end(c: &[char], i: usize) -> Option<usize> {
    let mut k = i + usize::from(c[i] == 'b');
    if c.get(k) != Some(&'r') {
        return None;
    }
    k += 1;
    let hashes = c[k..].iter().take_while(|&&h| h == '#').count();
    k += hashes;
    if c.get(k) != Some(&'"') {
        return None;
    }
    k += 1;
    while k < c.len() {
        if c[k] == '"'
            && c[k + 1..]
                .iter()
                .take(hashes)
                .filter(|&&h| h == '#')
                .count()
                == hashes
        {
            return Some(k + 1 + hashes);
        }
        k += 1;
    }
    Some(c.len())
}

/// A char literal, or `None` for a lifetime (`'a` with no closing quote).
fn char_end(c: &[char], i: usize) -> Option<usize> {
    match (c.get(i + 1), c.get(i + 2)) {
        (Some('\\'), _) => Some(quoted_end(c, i + 1, '\'')),
        (Some(_), Some('\'')) => Some(i + 3),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The walk itself drops a `#[cfg(test)] mod x;` file, not only the parts
    /// that find one. A fixture tree, because the wiring is the thing tested.
    #[test]
    fn a_test_only_module_file_is_not_production_code() {
        let root = std::env::temp_dir().join(format!("source_scan_{}", std::process::id()));
        let src = root.join("src");
        std::fs::create_dir_all(src.join("t")).unwrap();
        std::fs::write(src.join("lib.rs"), "mod p;\n#[cfg(test)]\nmod t;\n").unwrap();
        std::fs::write(src.join("p.rs"), "fn production() {}\n").unwrap();
        std::fs::write(src.join("t.rs"), "fn test_only() {}\nmod deeper;\n").unwrap();
        std::fs::write(src.join("t").join("deeper.rs"), "fn test_too() {}\n").unwrap();
        let mut found: Vec<String> = production_code(std::slice::from_ref(&src))
            .into_iter()
            .map(|(file, _)| file.strip_prefix(&src).unwrap().display().to_string())
            .collect();
        found.sort();
        std::fs::remove_dir_all(&root).unwrap();
        assert_eq!(found, vec!["lib.rs", "p.rs"]);
    }

    #[test]
    fn test_only_items_and_modules_are_not_production() {
        let text = concat!(
            "fn kept() {}\n",
            "#[cfg(test)]\n",
            "mod inline {\n",
            "    fn gone_inline() {}\n",
            "}\n",
            "#[cfg(test)]\n",
            "mod out_of_line;\n",
            "#[cfg(any(test, feature = \"x\"))]\n",
            "pub mod gated;\n",
            "#[cfg(not(test))]\n",
            "fn production_only() {}\n",
            "fn also_kept() {}\n",
        );
        let code = code_only(text);
        let stripped = strip_test_items(&code);
        for kept in ["fn kept", "fn production_only", "fn also_kept"] {
            assert!(stripped.contains(kept), "{kept} was dropped: {stripped}");
        }
        assert!(!stripped.contains("gone_inline"), "{stripped}");
        let lib = Path::new("/c/src/lib.rs");
        assert_eq!(
            test_module_paths(lib, &code),
            vec![
                PathBuf::from("/c/src/out_of_line"),
                PathBuf::from("/c/src/gated")
            ]
        );
        let nested = Path::new("/c/src/ui/reduce.rs");
        assert_eq!(
            test_module_paths(nested, &code)[0],
            PathBuf::from("/c/src/ui/reduce/out_of_line")
        );
        assert!(is_module_file(
            Path::new("/c/src/gated"),
            Path::new("/c/src/gated.rs")
        ));
        assert!(is_module_file(
            Path::new("/c/src/gated"),
            Path::new("/c/src/gated/a.rs")
        ));
        assert!(!is_module_file(
            Path::new("/c/src/gated"),
            Path::new("/c/src/gatedx.rs")
        ));
    }

    #[test]
    fn comments_and_literals_are_blanked_and_code_is_kept() {
        let text = concat!(
            "let a = Command::Kept(1); // Command::LineComment(2)\n",
            "/* Command::Block( /* nested */ Command::StillBlock ) */\n",
            "let s = \"Command::InString(\\\"x\\\")\";\n",
            "let r = r#\"Command::InRaw(\"q\")\"#;\n",
            "let q = '\"'; let l: &'static str = \"\"; Command::AfterChar(3)\n",
        );
        let kept = code_only(text);
        assert!(kept.contains("Command::Kept(1)"));
        assert!(kept.contains("Command::AfterChar(3)"));
        assert!(
            kept.contains("&'static str"),
            "a lifetime is not a char literal"
        );
        for gone in ["LineComment", "Block", "StillBlock", "InString", "InRaw"] {
            assert!(!kept.contains(gone), "{gone} survived: {kept}");
        }
        assert_eq!(kept.lines().count(), text.lines().count());
    }
}