Skip to main content

nix_compat/aterm/
escape.rs

1use std::sync::LazyLock;
2
3use aho_corasick::AhoCorasick;
4
5const PATTERNS: [&str; 5] = ["\\", "\n", "\r", "\t", "\""];
6const REPLACEMENTS: [&str; 5] = ["\\\\", "\\n", "\\r", "\\t", "\\\""];
7static AC: LazyLock<AhoCorasick> = LazyLock::new(|| {
8    AhoCorasick::builder()
9        .build(PATTERNS)
10        .expect("to init aho-corasick with PATTERNS")
11});
12
13/// Given a byte sequence, writes it in escaped form to the passed writer.
14/// Does not add surrounding quotes.
15pub fn write_escaped<P: AsRef<[u8]>>(s: P, w: &mut impl std::io::Write) -> std::io::Result<()> {
16    let s = s.as_ref();
17    let mut pos = 0;
18
19    for m in AC.find_iter(s) {
20        w.write_all(&s[pos..m.start()])?;
21        w.write_all(REPLACEMENTS[m.pattern().as_usize()].as_bytes())?;
22        pos = m.end();
23    }
24
25    w.write_all(&s[pos..])
26}
27
28#[cfg(test)]
29mod tests {
30    use super::write_escaped;
31    use rstest::rstest;
32
33    #[rstest]
34    #[case::empty(b"", b"")]
35    #[case::doublequote(b"\"", b"\\\"")]
36    #[case::colon(b":", b":")]
37    #[case::complex(b"foo\n\rbar\\baz", b"foo\\n\\rbar\\\\baz")]
38    fn escape(#[case] input: &[u8], #[case] expected: &[u8]) {
39        let mut buf = Vec::new();
40        write_escaped(input, &mut buf).unwrap();
41
42        assert_eq!(expected, buf.as_slice());
43    }
44}