qail-core 3.0.1

QAIL AST Kernel - typed query AST, schema validation, RLS context, and native access policy
//! Shared Rust lexical masking helpers used by build-time analyzers.

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
}

/// A `'` opens a char literal only as `'x'` (one UTF-8 char) or an escape
/// (`'\n'`, `'\''`, `'\u{1F600}'`). Anything else is a lifetime or a loop label
/// (`'a`, `'_`, `'static`, `'outer:`), which is code: read as a literal, it
/// blanked every line up to the next apostrophe, one inside a later comment
/// included, and the analyzers never saw the queries in between.
fn consume_char_literal(bytes: &[u8], start: usize) -> Option<usize> {
    let body = start + 1;
    let first = *bytes.get(body)?;
    if first == b'\\' {
        // `'\u{10FFFF}'` is the longest escape: its closing quote sits ten
        // bytes past the backslash.
        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}");
    }
}