3 files changed, 49 insertions, 17 deletions

diff --git a/crates/libeditor/src/completion.rs b/crates/libeditor/src/completion.rs
index 6335dba17..ec9388ee3 100644
--- a/crates/libeditor/src/completion.rs
+++ b/crates/libeditor/src/completion.rs
@@ -5,7 +5,6 @@ use libsyntax2::{
     ast::{self, NameOwner},
     algo::{
         ancestors,
-        visit::{visitor_ctx, VisitorCtx},
         walk::preorder,
         generate,
     },
@@ -52,24 +51,24 @@ fn compute_scopes(fn_def: ast::FnDef) -> FnScopes {
 
     let mut scope = root;
     if let Some(body) = fn_def.body() {
-        for child in body.syntax().children() {
-            let _ = visitor_ctx((&mut scopes, &mut scope))
-                .visit::<ast::LetStmt, _>(|stmt, (scopes, scope)| {
-                    *scope = scopes.new_scope(*scope);
+        for stmt in body.statements() {
+            match stmt {
+                ast::Stmt::LetStmt(stmt) => {
+                    scope = scopes.new_scope(scope);
                     if let Some(pat) = stmt.pat() {
-                        scopes.add_bindings(*scope, pat);
+                        scopes.add_bindings(scope, pat);
                     }
                     if let Some(expr) = stmt.initializer() {
-                        scopes.set_scope(expr.syntax(), *scope)
+                        scopes.set_scope(expr.syntax(), scope)
                     }
-                })
-                .visit::<ast::ExprStmt, _>(|expr, (scopes, scope)| {
-                    scopes.set_scope(expr.syntax(), *scope)
-                })
-                .visit::<ast::Expr, _>(|expr, (scopes, scope)| {
-                    scopes.set_scope(expr.syntax(), *scope)
-                })
-                .accept(child);
+                }
+                ast::Stmt::ExprStmt(expr) => {
+                    scopes.set_scope(expr.syntax(), scope)
+                }
+            }
+        }
+        if let Some(expr) = body.expr() {
+            scopes.set_scope(expr.syntax(), scope)
         }
     }
     scopes
diff --git a/crates/libsyntax2/src/ast/generated.rs b/crates/libsyntax2/src/ast/generated.rs
index b937fe5a2..1dd161f52 100644
--- a/crates/libsyntax2/src/ast/generated.rs
+++ b/crates/libsyntax2/src/ast/generated.rs
@@ -116,9 +116,13 @@ impl<'a> AstNode<'a> for Block<'a> {
 }
 
 impl<'a> Block<'a> {
-    pub fn let_stmts(self) -> impl Iterator<Item = LetStmt<'a>> + 'a {
+    pub fn statements(self) -> impl Iterator<Item = Stmt<'a>> + 'a {
         super::children(self)
     }
+
+    pub fn expr(self) -> Option<Expr<'a>> {
+        super::child_opt(self)
+    }
 }
 
 // BlockExpr
@@ -1370,6 +1374,31 @@ impl<'a> ast::TypeParamsOwner<'a> for StaticDef<'a> {}
 impl<'a> ast::AttrsOwner<'a> for StaticDef<'a> {}
 impl<'a> StaticDef<'a> {}
 
+// Stmt
+#[derive(Debug, Clone, Copy)]
+pub enum Stmt<'a> {
+    ExprStmt(ExprStmt<'a>),
+    LetStmt(LetStmt<'a>),
+}
+
+impl<'a> AstNode<'a> for Stmt<'a> {
+    fn cast(syntax: SyntaxNodeRef<'a>) -> Option<Self> {
+        match syntax.kind() {
+            EXPR_STMT => Some(Stmt::ExprStmt(ExprStmt { syntax })),
+            LET_STMT => Some(Stmt::LetStmt(LetStmt { syntax })),
+            _ => None,
+        }
+    }
+    fn syntax(self) -> SyntaxNodeRef<'a> {
+        match self {
+            Stmt::ExprStmt(inner) => inner.syntax(),
+            Stmt::LetStmt(inner) => inner.syntax(),
+        }
+    }
+}
+
+impl<'a> Stmt<'a> {}
+
 // StructDef
 #[derive(Debug, Clone, Copy)]
 pub struct StructDef<'a> {
diff --git a/crates/libsyntax2/src/grammar.ron b/crates/libsyntax2/src/grammar.ron
index aa2742b3e..f3c3d3036 100644
--- a/crates/libsyntax2/src/grammar.ron
+++ b/crates/libsyntax2/src/grammar.ron
@@ -437,9 +437,13 @@ Grammar(
             ["pat", "Pat"],
             ["initializer", "Expr"],
         ]),
+        "Stmt": (
+            enum: ["ExprStmt", "LetStmt"],
+        ),
         "Block": (
+            options: [["expr", "Expr"]],
             collections: [
-                ["let_stmts", "LetStmt"],
+                ["statements", "Stmt"],
             ]
         ),
         "ParamList": (