Parse string literals to find their length

src/lib.rs

@@ -274,13 +274,87 @@ impl Expect {
             line_start += line.len();
         }
         let (literal_start, line_indent) = target_line.unwrap();
-        let literal_length =
-            file[literal_start..].find("]]").expect("Couldn't find matching `]]` for `expect![[`.");
-        let literal_range = literal_start..literal_start + literal_length;
+
+        let lit_to_eof = &file[literal_start..];
+        let lit_to_eof_trimmed = lit_to_eof.trim_start();
+
+        let literal_start = literal_start + (lit_to_eof.len() - lit_to_eof_trimmed.len());
+
+        let literal_len = locate_end(lit_to_eof_trimmed)
+            .expect("Couldn't find matching `]]` for `expect![[`.");
+        let literal_range = literal_start..literal_start + literal_len;
         Location { line_indent, literal_range }
     }
 }
 
+fn locate_end(lit_to_eof: &str) -> Option<usize> {
+    assert!(lit_to_eof.chars().next().map_or(true, |c| !c.is_whitespace()));
+
+    if lit_to_eof.starts_with("]]") {
+        // expect![[ ]]
+        Some(0)
+    } else {
+        // expect![["foo"]]
+        find_str_lit_len(lit_to_eof)
+    }
+}
+
+/// Parses a string literal, returning the byte index of its last character
+/// (either a quote or a hash).
+fn find_str_lit_len(str_lit_to_eof: &str) -> Option<usize> {
+    use StrLitKind::*;
+    #[derive(Clone, Copy)]
+    enum StrLitKind {
+        Normal,
+        Raw(usize),
+    }
+
+    fn try_find_n_hashes(
+        s: &mut impl Iterator<Item = char>,
+        desired_hashes: usize,
+    ) -> Option<(usize, Option<char>)> {
+        let mut n = 0;
+        loop {
+            match s.next()? {
+                '#' => n += 1,
+                c => return Some((n, Some(c))),
+            }
+
+            if n == desired_hashes {
+                return Some((n, None));
+            }
+        }
+    }
+
+    let mut s = str_lit_to_eof.chars();
+    let kind = match s.next()? {
+        '"' => Normal,
+        'r' => {
+            let (n, c) = try_find_n_hashes(&mut s, usize::MAX)?;
+            if c != Some('"') { return None; }
+            Raw(n)
+        }
+        _ => return None,
+    };
+
+    let mut oldc = None;
+    loop {
+        let c = oldc.take().or_else(|| s.next())?;
+        match (c, kind) {
+            ('\\', Normal) => { let _escaped = s.next()?; }
+            ('"', Normal | Raw(0)) => break,
+            ('"', Raw(n)) => {
+                let (seen, c) = try_find_n_hashes(&mut s, n)?;
+                if seen == n { break; }
+                oldc = c;
+            }
+            _ => {}
+        }
+    }
+
+    Some(str_lit_to_eof.len() - s.as_str().len())
+}
+
 impl ExpectFile {
     /// Checks if file contents is equal to `actual`.
     pub fn assert_eq(&self, actual: &str) {
@@ -661,4 +735,48 @@ line1
 "#]],
         );
     }
+
+    #[test]
+    fn test_locate() {
+        macro_rules! check_locate {
+            ($( [[$s:literal]] ),* $(,)?) => {$({
+                let lit = stringify!($s);
+                let with_trailer = format!("{} \t]]\n", lit);
+                assert_eq!(locate_end(&with_trailer), Some(lit.len()));
+            })*};
+        }
+
+        // Check that we handle string literals containing "]]" correctly.
+        check_locate!(
+            [[r#"{ arr: [[1, 2], [3, 4]], other: "foo" } "#]],
+            [["]]"]],
+            [["\"]]"]],
+            [[r#""]]"#]],
+        );
+
+        // Check `expect![[  ]]` as well.
+        assert_eq!(locate_end("]]"), Some(0));
+    }
+
+    #[test]
+    fn test_find_str_lit_len() {
+        macro_rules! check_str_lit_len {
+            ($( $s:literal ),* $(,)?) => {$({
+                let lit = stringify!($s);
+                assert_eq!(find_str_lit_len(lit), Some(lit.len()));
+            })*}
+        }
+
+        check_str_lit_len![
+            r##"foa\""#"##,
+            r##"
+
+                asdf][]]""""#
+            "##,
+            "",
+            "\"",
+            "\"\"",
+            "#\"#\"#",
+        ];
+    }
 }