5 files changed, 72 insertions, 38 deletions

diff --git a/crates/libeditor/src/completion.rs b/crates/libeditor/src/completion.rs
index f733cd2b2..2b6a1db86 100644
--- a/crates/libeditor/src/completion.rs
+++ b/crates/libeditor/src/completion.rs
@@ -1,9 +1,11 @@
 use libsyntax2::{
-    File, TextUnit, AstNode,
-    ast::self,
+    File, TextUnit, AstNode, SyntaxKind::*,
+    ast::{self, LoopBodyOwner},
     algo::{
         ancestors,
+        visit::{visitor, Visitor},
     },
+    text_utils::is_subrange,
 };
 
 use {
@@ -30,11 +32,9 @@ pub fn scope_completion(file: &File, offset: TextUnit) -> Option<Vec<CompletionI
 
     let mut res = Vec::new();
     if let Some(fn_def) = ancestors(name_ref.syntax()).filter_map(ast::FnDef::cast).next() {
-        complete_keywords(&file, Some(fn_def), name_ref, &mut res);
+        complete_expr_keywords(&file, fn_def, name_ref, &mut res);
         let scopes = FnScopes::new(fn_def);
         complete_fn(name_ref, &scopes, &mut res);
-    } else {
-        complete_keywords(&file, None, name_ref, &mut res);
     }
     if let Some(root) = ancestors(name_ref.syntax()).filter_map(ast::Root::cast).next() {
         let scope = ModuleScope::new(root);
@@ -58,7 +58,7 @@ fn is_single_segment(name_ref: ast::NameRef) -> bool {
     }
 }
 
-fn complete_keywords(file: &File, fn_def: Option<ast::FnDef>, name_ref: ast::NameRef, acc: &mut Vec<CompletionItem>) {
+fn complete_expr_keywords(file: &File, fn_def: ast::FnDef, name_ref: ast::NameRef, acc: &mut Vec<CompletionItem>) {
     acc.push(keyword("if", "if $0 { }"));
     acc.push(keyword("match", "match $0 { }"));
     acc.push(keyword("while", "while $0 { }"));
@@ -72,22 +72,42 @@ fn complete_keywords(file: &File, fn_def: Option<ast::FnDef>, name_ref: ast::Nam
             }
         }
     }
+    if is_in_loop_body(name_ref) {
+        acc.push(keyword("continue", "continue"));
+        acc.push(keyword("break", "break"));
+    }
+    acc.extend(complete_return(fn_def, name_ref));
+}
 
-    if let Some(fn_def) = fn_def {
-        acc.extend(complete_return(fn_def, name_ref));
+fn is_in_loop_body(name_ref: ast::NameRef) -> bool {
+    for node in ancestors(name_ref.syntax()) {
+        if node.kind() == FN_DEF || node.kind() == LAMBDA_EXPR {
+            break;
+        }
+        let loop_body = visitor()
+            .visit::<ast::ForExpr, _>(LoopBodyOwner::loop_body)
+            .visit::<ast::WhileExpr, _>(LoopBodyOwner::loop_body)
+            .visit::<ast::LoopExpr, _>(LoopBodyOwner::loop_body)
+            .accept(node);
+        if let Some(Some(body)) = loop_body {
+            if is_subrange(body.syntax().range(), name_ref.syntax().range()) {
+                return true;
+            }
+        }
     }
+    false
 }
 
 fn complete_return(fn_def: ast::FnDef, name_ref: ast::NameRef) -> Option<CompletionItem> {
-    let is_last_in_block = ancestors(name_ref.syntax()).filter_map(ast::Expr::cast)
-        .next()
-        .and_then(|it| it.syntax().parent())
-        .and_then(ast::Block::cast)
-        .is_some();
+    // let is_last_in_block = ancestors(name_ref.syntax()).filter_map(ast::Expr::cast)
+    //     .next()
+    //     .and_then(|it| it.syntax().parent())
+    //     .and_then(ast::Block::cast)
+    //     .is_some();
 
-    if is_last_in_block {
-        return None;
-    }
+    // if is_last_in_block {
+    //     return None;
+    // }
 
     let is_stmt = match ancestors(name_ref.syntax()).filter_map(ast::ExprStmt::cast).next() {
         None => false,
@@ -220,7 +240,8 @@ mod tests {
             ", r#"[CompletionItem { name: "if", snippet: Some("if $0 { }") },
                    CompletionItem { name: "match", snippet: Some("match $0 { }") },
                    CompletionItem { name: "while", snippet: Some("while $0 { }") },
-                   CompletionItem { name: "loop", snippet: Some("loop {$0}") }]"#);
+                   CompletionItem { name: "loop", snippet: Some("loop {$0}") },
+                   CompletionItem { name: "return", snippet: Some("return") }]"#);
     }
 
     #[test]
@@ -236,7 +257,8 @@ mod tests {
                    CompletionItem { name: "while", snippet: Some("while $0 { }") },
                    CompletionItem { name: "loop", snippet: Some("loop {$0}") },
                    CompletionItem { name: "else", snippet: Some("else {$0}") },
-                   CompletionItem { name: "else if", snippet: Some("else if $0 { }") }]"#);
+                   CompletionItem { name: "else if", snippet: Some("else if $0 { }") },
+                   CompletionItem { name: "return", snippet: Some("return") }]"#);
     }
 
     #[test]
@@ -277,4 +299,19 @@ mod tests {
                    CompletionItem { name: "loop", snippet: Some("loop {$0}") },
                    CompletionItem { name: "return", snippet: Some("return $0") }]"#);
     }
+
+    #[test]
+    fn test_continue_break_completion() {
+        check_snippet_completion(r"
+            fn quux() -> i32 {
+                loop { <|> }
+            }
+            ", r#"[CompletionItem { name: "if", snippet: Some("if $0 { }") },
+                   CompletionItem { name: "match", snippet: Some("match $0 { }") },
+                   CompletionItem { name: "while", snippet: Some("while $0 { }") },
+                   CompletionItem { name: "loop", snippet: Some("loop {$0}") },
+                   CompletionItem { name: "continue", snippet: Some("continue") },
+                   CompletionItem { name: "break", snippet: Some("break") },
+                   CompletionItem { name: "return", snippet: Some("return $0") }]"#);
+    }
 }
diff --git a/crates/libeditor/src/scope/fn_scope.rs b/crates/libeditor/src/scope/fn_scope.rs
index de38b33f0..4b643237f 100644
--- a/crates/libeditor/src/scope/fn_scope.rs
+++ b/crates/libeditor/src/scope/fn_scope.rs
@@ -5,7 +5,7 @@ use std::{
 
 use libsyntax2::{
     SyntaxNodeRef, SyntaxNode, SmolStr, AstNode,
-    ast::{self, NameOwner},
+    ast::{self, NameOwner, LoopBodyOwner},
     algo::{ancestors, generate, walk::preorder}
 };
 
@@ -144,7 +144,7 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
             let cond_scope = e.condition().and_then(|cond| {
                 compute_cond_scopes(cond, scopes, scope)
             });
-            if let Some(block) = e.body() {
+            if let Some(block) = e.loop_body() {
                 compute_block_scopes(block, scopes, cond_scope.unwrap_or(scope));
             }
         },
@@ -162,7 +162,7 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
                 scope = scopes.new_scope(scope);
                 scopes.add_bindings(scope, pat);
             }
-            if let Some(block) = e.body() {
+            if let Some(block) = e.loop_body() {
                 compute_block_scopes(block, scopes, scope);
             }
         },
@@ -181,7 +181,7 @@ fn compute_expr_scopes(expr: ast::Expr, scopes: &mut FnScopes, scope: ScopeId) {
                 .for_each(|expr| compute_expr_scopes(expr, scopes, scope));
         }
         ast::Expr::LoopExpr(e) => {
-            if let Some(block) = e.body() {
+            if let Some(block) = e.loop_body() {
                 compute_block_scopes(block, scopes, scope);
             }
         }
diff --git a/crates/libsyntax2/src/ast/generated.rs b/crates/libsyntax2/src/ast/generated.rs
index 1bb3b11d1..0668dbfa7 100644
--- a/crates/libsyntax2/src/ast/generated.rs
+++ b/crates/libsyntax2/src/ast/generated.rs
@@ -593,6 +593,7 @@ impl<'a> AstNode<'a> for ForExpr<'a> {
     fn syntax(self) -> SyntaxNodeRef<'a> { self.syntax }
 }
 
+impl<'a> ast::LoopBodyOwner<'a> for ForExpr<'a> {}
 impl<'a> ForExpr<'a> {
     pub fn pat(self) -> Option<Pat<'a>> {
         super::child_opt(self)
@@ -601,10 +602,6 @@ impl<'a> ForExpr<'a> {
     pub fn iterable(self) -> Option<Expr<'a>> {
         super::child_opt(self)
     }
-
-    pub fn body(self) -> Option<Block<'a>> {
-        super::child_opt(self)
-    }
 }
 
 // ForType
@@ -845,11 +842,8 @@ impl<'a> AstNode<'a> for LoopExpr<'a> {
     fn syntax(self) -> SyntaxNodeRef<'a> { self.syntax }
 }
 
-impl<'a> LoopExpr<'a> {
-    pub fn body(self) -> Option<Block<'a>> {
-        super::child_opt(self)
-    }
-}
+impl<'a> ast::LoopBodyOwner<'a> for LoopExpr<'a> {}
+impl<'a> LoopExpr<'a> {}
 
 // MatchArm
 #[derive(Debug, Clone, Copy)]
@@ -2106,13 +2100,10 @@ impl<'a> AstNode<'a> for WhileExpr<'a> {
     fn syntax(self) -> SyntaxNodeRef<'a> { self.syntax }
 }
 
+impl<'a> ast::LoopBodyOwner<'a> for WhileExpr<'a> {}
 impl<'a> WhileExpr<'a> {
     pub fn condition(self) -> Option<Condition<'a>> {
         super::child_opt(self)
     }
-
-    pub fn body(self) -> Option<Block<'a>> {
-        super::child_opt(self)
-    }
 }
 
diff --git a/crates/libsyntax2/src/ast/mod.rs b/crates/libsyntax2/src/ast/mod.rs
index 3b5e9269f..49e283f5e 100644
--- a/crates/libsyntax2/src/ast/mod.rs
+++ b/crates/libsyntax2/src/ast/mod.rs
@@ -20,6 +20,12 @@ pub trait NameOwner<'a>: AstNode<'a> {
     }
 }
 
+pub trait LoopBodyOwner<'a>: AstNode<'a> {
+    fn loop_body(self) -> Option<Block<'a>> {
+        child_opt(self)
+    }
+}
+
 pub trait TypeParamsOwner<'a>: AstNode<'a> {
     fn type_param_list(self) -> Option<TypeParamList<'a>> {
         child_opt(self)
diff --git a/crates/libsyntax2/src/grammar.ron b/crates/libsyntax2/src/grammar.ron
index 8267c2854..da18da8a9 100644
--- a/crates/libsyntax2/src/grammar.ron
+++ b/crates/libsyntax2/src/grammar.ron
@@ -354,19 +354,19 @@ Grammar(
             options: [ ["condition", "Condition"] ]
         ),
         "LoopExpr": (
-            options: [ ["body", "Block"] ]
+            traits: ["LoopBodyOwner"],
         ),
         "ForExpr": (
+            traits: ["LoopBodyOwner"],
             options: [
                 ["pat", "Pat"],
                 ["iterable", "Expr"],
-                ["body", "Block"] ,
             ]
         ),
         "WhileExpr": (
+            traits: ["LoopBodyOwner"],
             options: [
                 ["condition", "Condition"],
-                ["body", "Block"],
             ]
         ),
         "ContinueExpr": (),