2 files changed, 63 insertions, 30 deletions

diff --git a/crates/ra_editor/src/folding_ranges.rs b/crates/ra_editor/src/folding_ranges.rs
index 8aabb5e0b..a1699d449 100644
--- a/crates/ra_editor/src/folding_ranges.rs
+++ b/crates/ra_editor/src/folding_ranges.rs
@@ -120,54 +120,65 @@ fn contiguous_range_for_comment<'a>(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use test_utils::extract_ranges;
+
+    fn do_check(text: &str, fold_kinds: &[FoldKind]) {
+        let (ranges, text) = extract_ranges(text);
+        let file = File::parse(&text);
+        let folds = folding_ranges(&file);
+
+        assert_eq!(folds.len(), ranges.len());
+        for ((fold, range), fold_kind) in folds.into_iter().zip(ranges.into_iter()).zip(fold_kinds.into_iter()) {
+            assert_eq!(fold.range.start(), range.start());
+            assert_eq!(fold.range.end(), range.end());
+            assert_eq!(&fold.kind, fold_kind);
+        }
+    }
 
     #[test]
     fn test_fold_comments() {
         let text = r#"
-// Hello
+<|>// Hello
 // this is a multiline
 // comment
-//
+//<|>
 
 // But this is not
 
 fn main() {
-    // We should
+    <|>// We should
     // also
     // fold
-    // this one.
+    // this one.<|>
+    <|>//! But this one is different
+    //! because it has another flavor<|>
+    <|>/* As does this
+    multiline comment */<|>
 }"#;
 
-        let file = File::parse(&text);
-        let folds = folding_ranges(&file);
-        assert_eq!(folds.len(), 2);
-        assert_eq!(folds[0].range.start(), 1.into());
-        assert_eq!(folds[0].range.end(), 46.into());
-        assert_eq!(folds[0].kind, FoldKind::Comment);
-
-        assert_eq!(folds[1].range.start(), 84.into());
-        assert_eq!(folds[1].range.end(), 137.into());
-        assert_eq!(folds[1].kind, FoldKind::Comment);
+        let fold_kinds = &[
+            FoldKind::Comment,
+            FoldKind::Comment,
+            FoldKind::Comment,
+            FoldKind::Comment,
+        ];
+        do_check(text, fold_kinds);
     }
 
     #[test]
     fn test_fold_imports() {
         let text = r#"
-use std::{
+<|>use std::{
     str,
     vec,
     io as iop
-};
+};<|>
 
 fn main() {
 }"#;
 
-        let file = File::parse(&text);
-        let folds = folding_ranges(&file);
-        assert_eq!(folds.len(), 1);
-        assert_eq!(folds[0].range.start(), 1.into());
-        assert_eq!(folds[0].range.end(), 46.into());
-        assert_eq!(folds[0].kind, FoldKind::Imports);
+        let folds = &[FoldKind::Imports];
+        do_check(text, folds);
     }
 
 
diff --git a/crates/test_utils/src/lib.rs b/crates/test_utils/src/lib.rs
index 068eb80ce..ee73153f0 100644
--- a/crates/test_utils/src/lib.rs
+++ b/crates/test_utils/src/lib.rs
@@ -38,22 +38,44 @@ pub fn assert_eq_dbg(expected: &str, actual: &impl fmt::Debug) {
 }
 
 pub fn extract_offset(text: &str) -> (TextUnit, String) {
-    let cursor = "<|>";
-    let cursor_pos = match text.find(cursor) {
+    match try_extract_offset(text) {
         None => panic!("text should contain cursor marker"),
-        Some(pos) => pos,
-    };
+        Some(result) => result,
+    }
+}
+
+pub fn try_extract_offset(text: &str) -> Option<(TextUnit, String)> {
+    let cursor = "<|>";
+    let cursor_pos = text.find(cursor)?;
     let mut new_text = String::with_capacity(text.len() - cursor.len());
     new_text.push_str(&text[..cursor_pos]);
     new_text.push_str(&text[cursor_pos + cursor.len()..]);
     let cursor_pos = TextUnit::from(cursor_pos as u32);
-    (cursor_pos, new_text)
+    Some((cursor_pos, new_text))
 }
 
 pub fn extract_range(text: &str) -> (TextRange, String) {
-    let (start, text) = extract_offset(text);
-    let (end, text) = extract_offset(&text);
-    (TextRange::from_to(start, end), text)
+    match try_extract_range(text) {
+        None => panic!("text should contain cursor marker"),
+        Some(result) => result,
+    }
+}
+
+pub fn try_extract_range(text: &str) -> Option<(TextRange, String)> {
+    let (start, text) = try_extract_offset(text)?;
+    let (end, text) = try_extract_offset(&text)?;
+    Some((TextRange::from_to(start, end), text))
+}
+
+pub fn extract_ranges(text: &str) -> (Vec<TextRange>, String) {
+    let mut ranges = Vec::new();
+    let mut text = String::from(text);
+    while let Some((range, new_text)) = try_extract_range(&text) {
+        text = new_text;
+        ranges.push(range);
+    }
+
+    (ranges, text)
 }
 
 pub fn add_cursor(text: &str, offset: TextUnit) -> String {