better extend selection

crates/libeditor/src/extend_selection.rs

@@ -1,7 +1,7 @@
 use libsyntax2::{
     File, TextRange, SyntaxNodeRef, TextUnit,
     SyntaxKind::*,
-    algo::{find_leaf_at_offset, find_covering_node, ancestors, Direction, siblings},
+    algo::{find_leaf_at_offset, LeafAtOffset, find_covering_node, ancestors, Direction, siblings},
 };
 
 pub fn extend_selection(file: &File, range: TextRange) -> Option<TextRange> {
@@ -13,10 +13,30 @@ pub(crate) fn extend(root: SyntaxNodeRef, range: TextRange) -> Option<TextRange>
     if range.is_empty() {
         let offset = range.start();
         let mut leaves = find_leaf_at_offset(root, offset);
-        if let Some(leaf) = leaves.clone().find(|node| node.kind() != WHITESPACE) {
+        if leaves.clone().all(|it| it.kind() == WHITESPACE) {
-            return Some(leaf.range());
+            return Some(extend_ws(root, leaves.next()?, offset));
         }
-        let ws = leaves.next()?;
+        let leaf = match leaves {
+            LeafAtOffset::None => return None,
+            LeafAtOffset::Single(l) => l,
+            LeafAtOffset::Between(l, r) => pick_best(l, r),
+        };
+        return Some(leaf.range());
+    };
+    let node = find_covering_node(root, range);
+    if node.kind() == COMMENT && range == node.range() {
+        if let Some(range) = extend_comments(node) {
+            return Some(range);
+        }
+    }
+
+    match ancestors(node).skip_while(|n| n.range() == range).next() {
+        None => None,
+        Some(parent) => Some(parent.range()),
+    }
+}
+
+fn extend_ws(root: SyntaxNodeRef, ws: SyntaxNodeRef, offset: TextUnit) -> TextRange {
     let ws_text = ws.leaf_text().unwrap();
     let suffix = TextRange::from_to(offset, ws.range().end()) - ws.range().start();
     let prefix = TextRange::from_to(ws.range().start(), offset) - ws.range().start();
@@ -33,21 +53,20 @@ pub(crate) fn extend(root: SyntaxNodeRef, range: TextRange) -> Option<TextRange>
             } else {
                 node.range().end()
             };
-                return Some(TextRange::from_to(start, end));
+            return TextRange::from_to(start, end);
-            }
-        }
-        return Some(ws.range());
-    };
-    let node = find_covering_node(root, range);
-    if node.kind() == COMMENT && range == node.range() {
-        if let Some(range) = extend_comments(node) {
-            return Some(range);
         }
     }
+    ws.range()
+}
 
-    match ancestors(node).skip_while(|n| n.range() == range).next() {
+fn pick_best<'a>(l: SyntaxNodeRef<'a>, r: SyntaxNodeRef<'a>) -> SyntaxNodeRef<'a> {
-        None => None,
+    return if priority(r) > priority(l) { r } else { l };
-        Some(parent) => Some(parent.range()),
+    fn priority(n: SyntaxNodeRef) -> usize {
+        match n.kind() {
+            WHITESPACE => 0,
+            IDENT | SELF_KW | SUPER_KW | CRATE_KW => 2,
+            _ => 1,
+        }
     }
 }
 
@@ -129,4 +148,20 @@ fn bar(){}
             &["// 1 + 1", "// fn foo() {\n// 1 + 1\n// }"]
         );
     }
+
+    #[test]
+    fn test_extend_selection_prefer_idents() {
+        do_check(
+            r#"
+fn main() { foo<|>+bar;}
+    "#,
+            &["foo", "foo+bar"]
+        );
+        do_check(
+            r#"
+fn main() { foo+<|>bar;}
+    "#,
+            &["bar", "foo+bar"]
+        );
+    }
 }