pub(super) fn starts_with_bytes(haystack: &[u8], idx: usize, needle: &[u8]) -> bool {
haystack
.get(idx..idx.saturating_add(needle.len()))
.is_some_and(|s| s == needle)
}
pub(super) fn consume_block_comment(bytes: &[u8], start: usize) -> usize {
let mut i = start + 2;
let mut depth = 1usize;
while i < bytes.len() && depth > 0 {
if starts_with_bytes(bytes, i, b"/*") {
depth += 1;
i += 2;
} else if starts_with_bytes(bytes, i, b"*/") {
depth = depth.saturating_sub(1);
i += 2;
} else {
i += 1;
}
}
i
}
pub(super) fn consume_rust_literal(bytes: &[u8], start: usize) -> Option<usize> {
if let Some((_, content_start, hashes)) = raw_string_prefix(bytes, start) {
let end_quote = find_raw_string_end(bytes, content_start, hashes)?;
return Some(end_quote + 1 + hashes);
}
if bytes.get(start).copied() == Some(b'"') || starts_with_bytes(bytes, start, b"b\"") {
let quote_offset = if bytes.get(start).copied() == Some(b'"') {
start
} else {
start + 1
};
let mut i = quote_offset + 1;
while i < bytes.len() {
if bytes[i] == b'\\' {
i = (i + 2).min(bytes.len());
continue;
}
if bytes[i] == b'"' {
return Some(i + 1);
}
i += 1;
}
return Some(bytes.len());
}
if bytes.get(start).copied() == Some(b'\'') {
return consume_char_literal(bytes, start);
}
None
}
fn consume_char_literal(bytes: &[u8], start: usize) -> Option<usize> {
let body = start + 1;
let first = *bytes.get(body)?;
if first == b'\\' {
let mut i = body + 2;
while i < bytes.len() && i <= body + 10 {
if bytes[i] == b'\'' {
return Some(i + 1);
}
i += 1;
}
return None;
}
let width = match first {
0xF0..=0xFF => 4,
0xE0..=0xEF => 3,
0xC0..=0xDF => 2,
_ => 1,
};
(bytes.get(body + width).copied() == Some(b'\'')).then_some(body + width + 1)
}
pub(super) fn mask_non_code(source: &str) -> String {
let bytes = source.as_bytes();
let mut out = bytes.to_vec();
let mut i = 0usize;
while i < bytes.len() {
if starts_with_bytes(bytes, i, b"//") {
let start = i;
i += 2;
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
for b in &mut out[start..i] {
*b = b' ';
}
continue;
}
if starts_with_bytes(bytes, i, b"/*") {
let start = i;
i = consume_block_comment(bytes, i);
for b in &mut out[start..i] {
if *b != b'\n' {
*b = b' ';
}
}
continue;
}
if let Some(next) = consume_rust_literal(bytes, i) {
for b in &mut out[i..next] {
if *b != b'\n' {
*b = b' ';
}
}
i = next;
continue;
}
i += 1;
}
String::from_utf8(out).unwrap_or_else(|_| source.to_string())
}
pub(super) fn raw_string_prefix(bytes: &[u8], idx: usize) -> Option<(usize, usize, usize)> {
if bytes.get(idx).copied() == Some(b'r') {
let mut j = idx + 1;
while bytes.get(j).copied() == Some(b'#') {
j += 1;
}
if bytes.get(j).copied() == Some(b'"') {
let hashes = j - (idx + 1);
return Some((idx, j + 1, hashes));
}
return None;
}
if bytes.get(idx).copied() == Some(b'b') && bytes.get(idx + 1).copied() == Some(b'r') {
let mut j = idx + 2;
while bytes.get(j).copied() == Some(b'#') {
j += 1;
}
if bytes.get(j).copied() == Some(b'"') {
let hashes = j - (idx + 2);
return Some((idx, j + 1, hashes));
}
}
None
}
pub(super) fn find_raw_string_end(bytes: &[u8], mut idx: usize, hashes: usize) -> Option<usize> {
while idx < bytes.len() {
if bytes[idx] == b'"' {
let mut ok = true;
for off in 0..hashes {
if bytes.get(idx + 1 + off).copied() != Some(b'#') {
ok = false;
break;
}
}
if ok {
return Some(idx);
}
}
idx += 1;
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn masks_comment_and_literal_content() {
let src = r#"
let _ = "query(SELECT 1)";
// comment with query_as!
/* block Qail::raw_sql */
let ok = 1;
"#;
let masked = mask_non_code(src);
assert!(!masked.contains("query(SELECT 1)"));
assert!(!masked.contains("query_as!"));
assert!(!masked.contains("Qail::raw_sql"));
assert!(masked.contains("let ok = 1;"));
}
#[test]
fn lifetimes_and_labels_do_not_mask_following_code() {
let src = r#"
fn day<'a>(scope: DayScope<'_>, name: &'a str) -> Qail {
'outer: loop {
conn.fetch_all(&cmd);
break 'outer;
}
}
// the caller's own orders
let visible = 1;
"#;
let masked = mask_non_code(src);
assert!(masked.contains("conn.fetch_all(&cmd);"), "{masked}");
assert!(masked.contains("'outer: loop {"), "{masked}");
assert!(masked.contains("let visible = 1;"), "{masked}");
assert!(!masked.contains("caller"), "{masked}");
}
#[test]
fn char_literals_stay_masked() {
let src = r#"let a = 'x'; let b = '\''; let c = '"'; let d = 'é'; let e = b'q'; let f = '\u{1F600}'; let tail = 1;"#;
let masked = mask_non_code(src);
for literal in ["'x'", r"'\''", "'\"'", "'é'", "'q'", r"'\u{1F600}'"] {
assert!(!masked.contains(literal), "{literal} visible in {masked}");
}
assert!(masked.contains("let tail = 1;"), "{masked}");
}
}