#[derive(Clone, Copy, PartialEq, Eq)]
enum Scan {
Code,
LineComment,
BlockComment(usize),
Str,
RawStr(usize),
CharLit,
}
#[must_use]
pub fn mask_non_code(source: &str) -> String {
let bytes = source.as_bytes();
let mut out = bytes.to_vec();
let mut state = Scan::Code;
let mut i = 0usize;
while i < bytes.len() {
let (next, advance) = match state {
Scan::Code => step_code(bytes, i, &mut out),
Scan::LineComment => step_line_comment(bytes, i, &mut out),
Scan::BlockComment(depth) => step_block_comment(bytes, i, depth, &mut out),
Scan::Str | Scan::CharLit => {
let terminator = if state == Scan::Str { b'"' } else { b'\'' };
step_quoted(bytes, i, terminator, state, &mut out)
}
Scan::RawStr(hashes) => step_raw_string(bytes, i, hashes, &mut out),
};
state = next;
i += advance;
}
String::from_utf8(out).unwrap_or_else(|_| source.to_owned())
}
const fn blank(out: &mut [u8], at: usize) {
if out[at] != b'\n' {
out[at] = b' ';
}
}
fn blank_run(out: &mut [u8], at: usize, count: usize) {
for offset in 0..count {
blank(out, at + offset);
}
}
fn step_code(bytes: &[u8], i: usize, out: &mut [u8]) -> (Scan, usize) {
if bytes[i] == b'/' && bytes.get(i + 1) == Some(&b'/') {
blank_run(out, i, 2);
(Scan::LineComment, 2)
} else if bytes[i] == b'/' && bytes.get(i + 1) == Some(&b'*') {
blank_run(out, i, 2);
(Scan::BlockComment(1), 2)
} else if bytes[i] == b'"' {
blank(out, i);
(Scan::Str, 1)
} else if let Some((after, prefix)) = raw_string_start(bytes, i) {
let hashes = after - i - prefix - 1;
blank_run(out, i, after - i);
(Scan::RawStr(hashes), after - i)
} else if bytes[i] == b'\'' && is_char_literal(bytes, i) {
blank(out, i);
(Scan::CharLit, 1)
} else {
(Scan::Code, 1)
}
}
const fn step_line_comment(bytes: &[u8], i: usize, out: &mut [u8]) -> (Scan, usize) {
if bytes[i] == b'\n' {
(Scan::Code, 1)
} else {
blank(out, i);
(Scan::LineComment, 1)
}
}
fn step_block_comment(bytes: &[u8], i: usize, depth: usize, out: &mut [u8]) -> (Scan, usize) {
if bytes[i] == b'/' && bytes.get(i + 1) == Some(&b'*') {
blank_run(out, i, 2);
(Scan::BlockComment(depth + 1), 2)
} else if bytes[i] == b'*' && bytes.get(i + 1) == Some(&b'/') {
blank_run(out, i, 2);
let next = if depth == 1 {
Scan::Code
} else {
Scan::BlockComment(depth - 1)
};
(next, 2)
} else {
blank(out, i);
(Scan::BlockComment(depth), 1)
}
}
fn step_quoted(
bytes: &[u8],
i: usize,
terminator: u8,
state: Scan,
out: &mut [u8],
) -> (Scan, usize) {
if bytes[i] == b'\\' {
let span = if i + 1 < bytes.len() { 2 } else { 1 };
blank_run(out, i, span);
(state, span)
} else if bytes[i] == terminator {
blank(out, i);
(Scan::Code, 1)
} else {
blank(out, i);
(state, 1)
}
}
fn step_raw_string(bytes: &[u8], i: usize, hashes: usize, out: &mut [u8]) -> (Scan, usize) {
if bytes[i] == b'"' && closes_raw_string(bytes, i, hashes) {
blank_run(out, i, hashes + 1);
(Scan::Code, hashes + 1)
} else {
blank(out, i);
(Scan::RawStr(hashes), 1)
}
}
fn raw_string_start(bytes: &[u8], i: usize) -> Option<(usize, usize)> {
let prefix = if bytes[i] == b'r' {
1
} else if (bytes[i] == b'b' || bytes[i] == b'c') && bytes.get(i + 1) == Some(&b'r') {
2
} else {
return None;
};
if i > 0 && (bytes[i - 1].is_ascii_alphanumeric() || bytes[i - 1] == b'_') {
return None;
}
let mut at = i + prefix;
while bytes.get(at) == Some(&b'#') {
at += 1;
}
(bytes.get(at) == Some(&b'"')).then_some((at + 1, prefix))
}
fn closes_raw_string(bytes: &[u8], i: usize, hashes: usize) -> bool {
(1..=hashes).all(|offset| bytes.get(i + offset) == Some(&b'#'))
}
fn is_char_literal(bytes: &[u8], i: usize) -> bool {
match bytes.get(i + 1) {
Some(b'\\') => true,
Some(_) => {
let mut at = i + 2;
while at < bytes.len() && at <= i + 5 {
if bytes[at] == b'\'' {
return true;
}
if bytes[at] & 0b1100_0000 != 0b1000_0000 {
return false;
}
at += 1;
}
false
}
None => false,
}
}
#[must_use]
pub fn balanced_close_paren(masked: &str, open: usize) -> Option<usize> {
let bytes = masked.as_bytes();
if bytes.get(open) != Some(&b'(') {
return None;
}
let mut depth = 0usize;
for (offset, byte) in bytes.iter().enumerate().skip(open) {
match byte {
b'(' => depth += 1,
b')' => {
depth -= 1;
if depth == 0 {
return Some(offset);
}
}
_ => {}
}
}
None
}
#[must_use]
pub fn code_lines(source: &str) -> Vec<(String, usize)> {
let masked = mask_non_code(source);
let mut out = Vec::new();
let mut offset = 0usize;
for line in masked.split_inclusive('\n') {
out.push((line.to_owned(), offset));
offset += line.len();
}
out
}
pub fn for_each_code_line<T>(
source: &str,
mut f: impl FnMut(&str, usize) -> Option<T>,
) -> Option<T> {
code_lines(source)
.into_iter()
.find_map(|(line, offset)| f(&line, offset))
}
#[must_use]
pub fn async_main_offset(line: &str) -> Option<usize> {
const NEEDLE: &str = "async fn main";
let mut base = 0;
while let Some(at) = line[base..].find(NEEDLE) {
let start = base + at;
let end = start + NEEDLE.len();
if line[end..]
.chars()
.next()
.is_none_or(|c| !c.is_alphanumeric() && c != '_')
{
return Some(start);
}
base = end;
}
None
}
#[cfg(test)]
mod tests {
use super::{
async_main_offset, balanced_close_paren, code_lines, for_each_code_line, mask_non_code,
};
#[test]
fn masking_preserves_length_and_newlines() {
for source in [
"let x = \"hello\";\n",
"// comment\nlet y = 2;\n",
"/* a\n b */\nlet z = 3;\n",
"let r = r#\"raw \" string\"#;\n",
"let c = 'x'; let s: &'a str;\n",
] {
let masked = mask_non_code(source);
assert_eq!(masked.len(), source.len(), "{source:?}");
assert_eq!(
masked.matches('\n').count(),
source.matches('\n').count(),
"{source:?}"
);
}
}
#[test]
fn string_contents_are_not_code() {
let source = "fn main() {\n let help = \"run AdminPlugin::new() yourself\";\n}\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("AdminPlugin::new(")
.then_some(()))
.is_none()
);
}
#[test]
fn raw_string_contents_are_not_code() {
let source = "fn main() {\n let doc = r#\"\n autumn_web::app()\n\"#;\n}\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_none()
);
}
#[test]
fn raw_byte_and_c_string_contents_are_not_code() {
for source in [
"fn main() {\n let doc = br#\"\n autumn_web::app()\n\"#;\n}\n",
"fn main() {\n let doc = cr##\"\n autumn_web::app()\n\"##;\n}\n",
] {
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_none(),
"{source:?}"
);
}
}
#[test]
fn a_quote_inside_a_raw_byte_string_does_not_desync_the_mask() {
let source = "fn main() {\n let doc = br#\"a \" b autumn_web::app()\"#;\n let real = autumn_web::app();\n}\n";
let found = for_each_code_line(source, |line, offset| {
line.contains("autumn_web::app()").then_some(offset)
})
.expect("the real call");
assert!(
source[found..].trim_start().starts_with("let real"),
"the probe leaked as code: {source:?}"
);
}
#[test]
fn a_prefix_continuing_an_identifier_is_not_a_raw_string() {
let source = "let xbr = 1;\nlet s = xbr\"not raw\";\nlet real = autumn_web::app();\n";
let found = for_each_code_line(source, |line, offset| {
line.contains("autumn_web::app()").then_some(offset)
})
.expect("the real call");
assert!(
source[found..].starts_with("let real"),
"the scan desynced: {source:?}"
);
}
#[test]
fn raw_strings_respect_their_hash_count() {
let source = "let d = r##\"a \"# b autumn_web::app()\"##;\nlet real = autumn_web::app();\n";
let found = for_each_code_line(source, |line, offset| {
line.contains("autumn_web::app()").then_some(offset)
})
.expect("the real call");
assert!(source[found..].starts_with("let real"), "{found}");
}
#[test]
fn line_and_block_comments_are_not_code() {
for source in [
"// let x = 1;\n",
"/// let x = 1;\n",
"//! let x = 1;\n",
" // let x = 1;\n",
"/*\nlet x = 1;\n*/\n",
"/* /* nested */ let x = 1; */\n",
] {
assert!(
for_each_code_line(source, |line, _| line.contains("let x").then_some(()))
.is_none(),
"{source:?}"
);
}
}
#[test]
fn comment_and_string_states_do_not_leak_into_each_other() {
let source = "let url = \"https://example.com\";\nlet real = autumn_web::app();\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_some()
);
let source = "// he said \"hi\n let real = autumn_web::app();\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_some(),
"an unbalanced quote in a comment must not swallow the rest of the file"
);
}
#[test]
fn lifetimes_do_not_open_a_character_literal() {
let source = "fn f<'a>(s: &'a str) {}\nlet real = autumn_web::app();\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_some()
);
}
#[test]
fn character_literals_are_still_masked() {
let source = "let q = '\"'; let real = autumn_web::app();\n";
assert!(
for_each_code_line(source, |line, _| line
.contains("autumn_web::app()")
.then_some(()))
.is_some(),
"a quote inside a char literal must not open a string"
);
}
#[test]
fn the_offset_locates_the_matching_line() {
let source = "let x = 1;\nlet x = 1;\nlet y = 2;\n";
let offset = for_each_code_line(source, |line, offset| {
line.contains("y = 2").then_some(offset)
})
.expect("found");
assert_eq!(&source[offset..offset + 9], "let y = 2");
}
#[test]
fn every_line_is_reported_with_a_stable_offset() {
let source = "a\nbb\nccc\n";
let lines = code_lines(source);
assert_eq!(lines.len(), 3);
assert_eq!(lines[0].1, 0);
assert_eq!(lines[1].1, 2);
assert_eq!(lines[2].1, 5);
}
#[test]
fn a_helper_whose_name_starts_with_main_is_not_the_entry_point() {
assert!(async_main_offset("async fn main() {").is_some());
assert!(async_main_offset("pub async fn main() {").is_some());
assert!(async_main_offset("async fn main_loop() {").is_none());
assert!(async_main_offset("async fn mainly() {").is_none());
assert!(async_main_offset("fn main() {").is_none());
}
#[test]
fn balanced_close_paren_spans_nested_calls() {
let src = "f(a, g(b, c), d) tail";
assert_eq!(balanced_close_paren(src, 1), Some(15));
assert_eq!(balanced_close_paren(src, 0), None, "index 0 is not a paren");
assert_eq!(balanced_close_paren("f(a", 1), None, "unterminated");
}
#[test]
fn the_first_match_wins() {
let source = "let a = 1;\nlet a = 2;\n";
let offset = for_each_code_line(source, |line, offset| {
line.contains("let a").then_some(offset)
})
.expect("found");
assert_eq!(offset, 0);
}
}