From cde2a1de3616d2e186555572996eecbd37e597e5 Mon Sep 17 00:00:00 2001
From: Jonas Schievink <jonas.schievink@ferrous-systems.com>
Date: Tue, 12 Apr 2022 19:39:19 +0200
Subject: [PATCH] Add trailing `;` when typing `=` in assignment

---
 crates/ide/src/typing.rs | 104 +++++++++++++++++++++++++++++++++++----
 1 file changed, 95 insertions(+), 9 deletions(-)

diff --git a/crates/ide/src/typing.rs b/crates/ide/src/typing.rs
index 0095012791..62c6ae5c1d 100644
--- a/crates/ide/src/typing.rs
+++ b/crates/ide/src/typing.rs
@@ -166,11 +166,34 @@ fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
     if !stdx::always!(file.syntax().text().char_at(offset) == Some('=')) {
         return None;
     }
-    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
-    if let_stmt.semicolon_token().is_some() {
-        return None;
+
+    if let Some(edit) = let_stmt(file, offset) {
+        return Some(edit);
     }
-    if let Some(expr) = let_stmt.initializer() {
+    if let Some(edit) = assign_expr(file, offset) {
+        return Some(edit);
+    }
+
+    return None;
+
+    fn assign_expr(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
+        let binop: ast::BinExpr = find_node_at_offset(file.syntax(), offset)?;
+        if !matches!(binop.op_kind(), Some(ast::BinaryOp::Assignment { op: None })) {
+            return None;
+        }
+
+        if let Some(expr_stmt) = ast::ExprStmt::cast(binop.syntax().parent()?) {
+            if expr_stmt.semicolon_token().is_some() {
+                return None;
+            }
+        } else {
+            if !ast::StmtList::can_cast(binop.syntax().parent()?.kind()) {
+                // Parent must be `ExprStmt` or `StmtList` for `;` to be valid.
+                return None;
+            }
+        }
+
+        let expr = binop.rhs()?;
         let expr_range = expr.syntax().text_range();
         if expr_range.contains(offset) && offset != expr_range.start() {
             return None;
@@ -178,11 +201,26 @@ fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
         if file.syntax().text().slice(offset..expr_range.start()).contains_char('\n') {
             return None;
         }
-    } else {
-        return None;
+        let offset = expr.syntax().text_range().end();
+        Some(TextEdit::insert(offset, ";".to_string()))
+    }
+
+    fn let_stmt(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
+        let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
+        if let_stmt.semicolon_token().is_some() {
+            return None;
+        }
+        let expr = let_stmt.initializer()?;
+        let expr_range = expr.syntax().text_range();
+        if expr_range.contains(offset) && offset != expr_range.start() {
+            return None;
+        }
+        if file.syntax().text().slice(offset..expr_range.start()).contains_char('\n') {
+            return None;
+        }
+        let offset = let_stmt.syntax().text_range().end();
+        Some(TextEdit::insert(offset, ";".to_string()))
     }
-    let offset = let_stmt.syntax().text_range().end();
-    Some(TextEdit::insert(offset, ";".to_string()))
 }
 
 /// Returns an edit which should be applied when a dot ('.') is typed on a blank line, indenting the line appropriately.
@@ -286,7 +324,7 @@ mod tests {
     }
 
     #[test]
-    fn test_on_eq_typed() {
+    fn test_semi_after_let() {
         //     do_check(r"
         // fn foo() {
         //     let foo =$0
@@ -322,6 +360,54 @@ fn foo() {
         // ");
     }
 
+    #[test]
+    fn test_semi_after_assign() {
+        type_char(
+            '=',
+            r#"
+fn f() {
+    i $0 0
+}
+"#,
+            r#"
+fn f() {
+    i = 0;
+}
+"#,
+        );
+        type_char(
+            '=',
+            r#"
+fn f() {
+    i $0 0
+    i
+}
+"#,
+            r#"
+fn f() {
+    i = 0;
+    i
+}
+"#,
+        );
+        type_char_noop(
+            '=',
+            r#"
+fn f(x: u8) {
+    if x $0
+}
+"#,
+        );
+        type_char_noop(
+            '=',
+            r#"
+fn f() {
+    g(i $0 0);
+}
+"#,
+        );
+    }
+
     #[test]
     fn indents_new_chain_call() {
         type_char(