tweak extend selection

3 files changed, 54 insertions, 16 deletions

diff --git a/crates/libeditor/src/extend_selection.rs b/crates/libeditor/src/extend_selection.rs
index 30cff6558..3adb1e45e 100644
--- a/crates/libeditor/src/extend_selection.rs
+++ b/crates/libeditor/src/extend_selection.rs
@@ -1,5 +1,5 @@
 use libsyntax2::{
-    File, TextRange, SyntaxNodeRef,
+    File, TextRange, SyntaxNodeRef, TextUnit,
     SyntaxKind::*,
     algo::{find_leaf_at_offset, find_covering_node, ancestors, Direction, siblings},
 };
@@ -18,11 +18,22 @@ pub(crate) fn extend(root: SyntaxNodeRef, range: TextRange) -> Option<TextRange>
         }
         let ws = leaves.next()?;
         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];
+        let suffix = TextRange::from_to(offset, ws.range().end()) - ws.range().start();
+        let prefix = TextRange::from_to(ws.range().start(), offset) - ws.range().start();
+        let ws_suffix = &ws_text.as_str()[suffix];
+        let ws_prefix = &ws_text.as_str()[prefix];
         if ws_text.contains("\n") && !ws_suffix.contains("\n") {
             if let Some(node) = ws.next_sibling() {
-                return Some(node.range());
+                let start = match ws_prefix.rfind('\n') {
+                    Some(idx) => ws.range().start() + TextUnit::from((idx + 1) as u32),
+                    None => node.range().start()
+                };
+                let end = if root.text().char_at(node.range().end()) == Some('\n') {
+                    node.range().end() + TextUnit::of_char('\n')
+                } else {
+                    node.range().end()
+                };
+                return Some(TextRange::from_to(start, end));
             }
         }
         return Some(ws.range());
@@ -99,7 +110,7 @@ impl S {
 
     }
 }"#,
-            &["fn foo() {\n\n    }"]
+            &["    fn foo() {\n\n    }\n"]
         );
     }
 
diff --git a/crates/libeditor/src/scope/fn_scope.rs b/crates/libeditor/src/scope/fn_scope.rs
index 4b643237f..5c04e2f9b 100644
--- a/crates/libeditor/src/scope/fn_scope.rs
+++ b/crates/libeditor/src/scope/fn_scope.rs
@@ -140,6 +140,16 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
                 compute_block_scopes(block, scopes, scope);
             }
         },
+        ast::Expr::BlockExpr(e) => {
+            if let Some(block) = e.block() {
+                compute_block_scopes(block, scopes, scope);
+            }
+        }
+        ast::Expr::LoopExpr(e) => {
+            if let Some(block) = e.loop_body() {
+                compute_block_scopes(block, scopes, scope);
+            }
+        }
         ast::Expr::WhileExpr(e) => {
             let cond_scope = e.condition().and_then(|cond| {
                 compute_cond_scopes(cond, scopes, scope)
@@ -147,11 +157,6 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
             if let Some(block) = e.loop_body() {
                 compute_block_scopes(block, scopes, cond_scope.unwrap_or(scope));
             }
-        },
-        ast::Expr::BlockExpr(e) => {
-            if let Some(block) = e.block() {
-                compute_block_scopes(block, scopes, scope);
-            }
         }
         ast::Expr::ForExpr(e) => {
             if let Some(expr) = e.iterable() {
@@ -165,7 +170,7 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
             if let Some(block) = e.loop_body() {
                 compute_block_scopes(block, scopes, scope);
             }
-        },
+        }
         ast::Expr::LambdaExpr(e) => {
             let mut scope = scopes.new_scope(scope);
             scopes.add_params_bindings(scope, e.param_list());
@@ -180,11 +185,7 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
                 .chain(e.expr())
                 .for_each(|expr| compute_expr_scopes(expr, scopes, scope));
         }
-        ast::Expr::LoopExpr(e) => {
-            if let Some(block) = e.loop_body() {
-                compute_block_scopes(block, scopes, scope);
-            }
-        }
+
         _ => {
             expr.syntax().children()
                 .filter_map(ast::Expr::cast)
@@ -273,4 +274,18 @@ mod tests {
             &["x"],
         );
     }
+
+    // #[test]
+    // fn test_match() {
+    //     do_check(r"
+    //         fn quux() {
+    //             match () {
+    //                 Some(x) => {
+    //                     <|>
+    //                 }
+    //             };
+    //         }",
+    //         &["x"],
+    //     );
+    // }
 }
diff --git a/crates/libsyntax2/src/yellow/syntax_text.rs b/crates/libsyntax2/src/yellow/syntax_text.rs
index b855687c4..280bedd78 100644
--- a/crates/libsyntax2/src/yellow/syntax_text.rs
+++ b/crates/libsyntax2/src/yellow/syntax_text.rs
@@ -61,6 +61,18 @@ impl<'a> SyntaxText<'a> {
             });
         SyntaxText { node: self.node, range }
     }
+    pub fn char_at(&self, offset: TextUnit) -> Option<char> {
+        let mut start: TextUnit = 0.into();
+        for chunk in self.chunks() {
+            let end = start + TextUnit::of_str(chunk);
+            if start <= offset && offset < end {
+                let off: usize = u32::from(offset - start) as usize;
+                return Some(chunk[off..].chars().next().unwrap());
+            }
+            start = end;
+        }
+        None
+    }
 }
 
 impl<'a> fmt::Debug for SyntaxText<'a> {