3 files changed, 27 insertions, 9 deletions
diff --git a/lib/syntax_suggest/around_block_scan.rb b/lib/syntax_suggest/around_block_scan.rb
index 21f309bf11..346a7a99ad 100644
--- a/lib/syntax_suggest/around_block_scan.rb
+++ b/lib/syntax_suggest/around_block_scan.rb
@@ -181,7 +181,7 @@ module SyntaxSuggest
           lines << line if line.is_kw? || line.is_end?
         }
       )
-      @scanner.try_rollback
+      @scanner.stash_changes
       lines
     end
 
@@ -240,7 +240,7 @@ module SyntaxSuggest
           true
         }
       )
-      @scanner.try_rollback
+      @scanner.stash_changes
       self
     end
 
@@ -275,26 +275,34 @@ module SyntaxSuggest
 
       return self if kw_count == end_count # nothing to balance
 
+      @scanner.commit_if_changed
+      scan_while { |line| line.empty? }
+
       # More ends than keywords, check if we can balance expanding up
       next_up = @scanner.next_up
       next_down = @scanner.next_down
       if (end_count - kw_count) == 1 && next_up
-        if next_up.is_kw? && next_up.indent >= @orig_indent
+        if next_up.is_kw? && next_up.indent >= @target_indent
           @scanner.scan(
             up: ->(line, _, _) { line == next_up },
             down: ->(line, _, _) { false }
           )
+          @scanner.commit_if_changed
         end
 
       # More keywords than ends, check if we can balance by expanding down
       elsif (kw_count - end_count) == 1 && next_down
-        if next_down.is_end? && next_down.indent >= @orig_indent
+        if next_down.is_end? && next_down.indent >= @target_indent
           @scanner.scan(
             up: ->(line, _, _) { false },
-            down: ->(line) { line == next_down }
+            down: ->(line, _, _) { line == next_down }
           )
+          @scanner.commit_if_changed
         end
       end
+
+      @scanner.stash_changes
+
       self
     end
 
diff --git a/lib/syntax_suggest/block_expand.rb b/lib/syntax_suggest/block_expand.rb
index e85d286768..e9b486c720 100644
--- a/lib/syntax_suggest/block_expand.rb
+++ b/lib/syntax_suggest/block_expand.rb
@@ -61,11 +61,14 @@ module SyntaxSuggest
     # they can expand to capture more code up and down). It does this conservatively
     # as there's no undo (currently).
     def expand_indent(block)
-      AroundBlockScan.new(code_lines: @code_lines, block: block)
+      now = AroundBlockScan.new(code_lines: @code_lines, block: block)
         .force_add_hidden
         .stop_after_kw
         .scan_adjacent_indent
-        .code_block
+
+      now.lookahead_balance_one_line
+
+      now.code_block
     end
 
     # A neighbor is code that is at or above the current indent line.
diff --git a/lib/syntax_suggest/scan_history.rb b/lib/syntax_suggest/scan_history.rb
index c5664d0f65..2be4aaf3c0 100644
--- a/lib/syntax_suggest/scan_history.rb
+++ b/lib/syntax_suggest/scan_history.rb
@@ -17,13 +17,20 @@ module SyntaxSuggest
     def commit_if_changed
       if changed?
         @history << CodeBlock.new(lines: @code_lines[before_index..after_index])
-        refresh_index
       end
 
       self
     end
 
-    def try_rollback
+    # Discards any changes that have not been committed
+    def stash_changes
+      refresh_index
+    end
+
+    # Discard changes that have not been committed and revert the last commit
+    #
+    # Cannot revert the first commit
+    def revert_last_commit
       if @history.length > 1
         @history.pop
         refresh_index