4 files changed, 53 insertions, 22 deletions

diff --git a/crates/libeditor/src/extend_selection.rs b/crates/libeditor/src/extend_selection.rs
index 171e40692..32873f491 100644
--- a/crates/libeditor/src/extend_selection.rs
+++ b/crates/libeditor/src/extend_selection.rs
@@ -17,18 +17,14 @@ pub(crate) fn extend(root: SyntaxNodeRef, range: TextRange) -> Option<TextRange>
             return Some(leaf.range());
         }
         let ws = leaves.next()?;
-//        let ws_suffix = file.text().slice(
-//            TextRange::from_to(offset, ws.range().end())
-//        );
-//        if ws.text().contains("\n") && !ws_suffix.contains("\n") {
-//            if let Some(line_end) = file.text()
-//                .slice(TextSuffix::from(ws.range().end()))
-//                .find("\n")
-//            {
-//                let range = TextRange::from_len(ws.range().end(), line_end);
-//                return Some(find_covering_node(file.root(), range).range());
-//            }
-//        }
+        let ws_text = ws.leaf_text().unwrap();
+        let range = TextRange::from_to(offset, ws.range().end()) - ws.range().start();
+        let ws_suffix = &ws_text.as_str()[range];
+        if ws_text.contains("\n") && !ws_suffix.contains("\n") {
+            if let Some(node) = ws.next_sibling() {
+                return Some(node.range());
+            }
+        }
         return Some(ws.range());
     };
     let node = find_covering_node(root, range);
diff --git a/crates/libeditor/tests/test.rs b/crates/libeditor/tests/test.rs
index a2e26003a..f8365949e 100644
--- a/crates/libeditor/tests/test.rs
+++ b/crates/libeditor/tests/test.rs
@@ -12,15 +12,32 @@ use libeditor::{
 
 #[test]
 fn test_extend_selection() {
-    let file = file(r#"fn foo() {
-    1 + 1
-}
-"#);
-    let range = TextRange::offset_len(18.into(), 0.into());
-    let range = extend_selection(&file, range).unwrap();
-    assert_eq!(range, TextRange::from_to(17.into(), 18.into()));
-    let range = extend_selection(&file, range).unwrap();
-    assert_eq!(range, TextRange::from_to(15.into(), 20.into()));
+    fn do_check(before: &str, afters: &[&str]) {
+        let (cursor, before) = extract_cursor(before);
+        let file = file(&before);
+        let mut range = TextRange::offset_len(cursor, 0.into());
+        for &after in afters {
+            range = extend_selection(&file, range)
+                .unwrap();
+            let actual = &before[range];
+            assert_eq!(after, actual);
+        }
+    }
+
+    do_check(
+        r#"fn foo() { <|>1 + 1 }"#,
+        &["1", "1 + 1", "{ 1 + 1 }"],
+    );
+
+    do_check(
+        r#"
+impl S {
+<|>    fn foo() {
+
+    }
+}"#,
+        &["fn foo() {\n\n    }"]
+    );
 }
 
 #[test]
diff --git a/crates/libsyntax2/Cargo.toml b/crates/libsyntax2/Cargo.toml
index b55cdf481..e8810046a 100644
--- a/crates/libsyntax2/Cargo.toml
+++ b/crates/libsyntax2/Cargo.toml
@@ -6,7 +6,7 @@ license = "MIT OR Apache-2.0"
 
 [dependencies]
 unicode-xid = "0.1.0"
-text_unit = "0.1.2"
+text_unit = "0.1.3"
 itertools = "0.7.8"
 drop_bomb = "0.1.4"
 parking_lot = "0.6.0"
diff --git a/crates/libsyntax2/src/ast/generated.rs b/crates/libsyntax2/src/ast/generated.rs
index 610b5198c..0651da26d 100644
--- a/crates/libsyntax2/src/ast/generated.rs
+++ b/crates/libsyntax2/src/ast/generated.rs
@@ -641,3 +641,21 @@ impl<'a> AstNode<'a> for TypeRef<'a> {
 
 impl<'a> TypeRef<'a> {}
 
+// Whitespace
+#[derive(Debug, Clone, Copy)]
+pub struct Whitespace<'a> {
+    syntax: SyntaxNodeRef<'a>,
+}
+
+impl<'a> AstNode<'a> for Whitespace<'a> {
+    fn cast(syntax: SyntaxNodeRef<'a>) -> Option<Self> {
+        match syntax.kind() {
+            WHITESPACE => Some(Whitespace { syntax }),
+            _ => None,
+        }
+    }
+    fn syntax(self) -> SyntaxNodeRef<'a> { self.syntax }
+}
+
+impl<'a> Whitespace<'a> {}
+