3 files changed, 140 insertions, 9 deletions
diff --git a/lib/syntax_suggest/around_block_scan.rb b/lib/syntax_suggest/around_block_scan.rb
index 2a57d1b19e..4793c3b5e6 100644
--- a/lib/syntax_suggest/around_block_scan.rb
+++ b/lib/syntax_suggest/around_block_scan.rb
@@ -61,7 +61,6 @@ module SyntaxSuggest
 
     def scan_while
       stop_next = false
-
       kw_count = 0
       end_count = 0
       index = before_lines.reverse_each.take_while do |line|
@@ -166,7 +165,55 @@ module SyntaxSuggest
       end
     end
 
-    def scan_neighbors
+    # Scanning is intentionally conservative because
+    # we have no way of rolling back an agressive block (at this time)
+    #
+    # If a block was stopped for some trivial reason, (like an empty line)
+    # but the next line would have caused it to be balanced then we
+    # can check that condition and grab just one more line either up or
+    # down.
+    #
+    # For example, below if we're scanning up, line 2 might cause
+    # the scanning to stop. This is because empty lines might
+    # denote logical breaks where the user intended to chunk code
+    # which is a good place to stop and check validity. Unfortunately
+    # it also means we might have a "dangling" keyword or end.
+    #
+    #   1 def bark
+    #   2
+    #   3 end
+    #
+    # If lines 2 and 3 are in the block, then when this method is
+    # run it would see it is unbalanced, but that acquiring line 1
+    # would make it balanced, so that's what it does.
+    def lookahead_balance_one_line
+      kw_count = 0
+      end_count = 0
+      lines.each do |line|
+        kw_count += 1 if line.is_kw?
+        end_count += 1 if line.is_end?
+      end
+
+      return self if kw_count == end_count # nothing to balance
+
+      # More ends than keywords, check if we can balance expanding up
+      if (end_count - kw_count) == 1 && next_up
+        return self unless next_up.is_kw?
+        return self unless next_up.indent >= @orig_indent
+
+        @before_index = next_up.index
+
+      # More keywords than ends, check if we can balance by expanding down
+      elsif (kw_count - end_count) == 1 && next_down
+        return self unless next_down.is_end?
+        return self unless next_down.indent >= @orig_indent
+
+        @after_index = next_down.index
+      end
+      self
+    end
+
+    def scan_neighbors_not_empty
       scan_while { |line| line.not_empty? && line.indent >= @orig_indent }
     end
 
diff --git a/lib/syntax_suggest/block_expand.rb b/lib/syntax_suggest/block_expand.rb
index 396b2c3a1a..8142e74869 100644
--- a/lib/syntax_suggest/block_expand.rb
+++ b/lib/syntax_suggest/block_expand.rb
@@ -35,14 +35,31 @@ module SyntaxSuggest
       @code_lines = code_lines
     end
 
+    # Main interface. Expand current indentation, before
+    # expanding to a lower indentation
     def call(block)
       if (next_block = expand_neighbors(block))
-        return next_block
+        next_block
+      else
+        expand_indent(block)
       end
-
-      expand_indent(block)
     end
 
+    # Expands code to the next lowest indentation
+    #
+    # For example:
+    #
+    #   1 def dog
+    #   2   print "dog"
+    #   3 end
+    #
+    # If a block starts on line 2 then it has captured all it's "neighbors" (code at
+    # the same indentation or higher). To continue expanding, this block must capture
+    # lines one and three which are at a different indentation level.
+    #
+    # This method allows fully expanded blocks to decrease their indentation level (so
+    # 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)
         .skip(:hidden?)
@@ -51,14 +68,82 @@ module SyntaxSuggest
         .code_block
     end
 
+    # A neighbor is code that is at or above the current indent line.
+    #
+    # First we build a block with all neighbors. If we can't go further
+    # then we decrease the indentation threshold and expand via indentation
+    # i.e. `expand_indent`
+    #
+    # Handles two general cases.
+    #
+    # ## Case #1: Check code inside of methods/classes/etc.
+    #
+    # It's important to note, that not everything in a given indentation level can be parsed
+    # as valid code even if it's part of valid code. For example:
+    #
+    #   1 hash = {
+    #   2   name: "richard",
+    #   3   dog: "cinco",
+    #   4 }
+    #
+    # In this case lines 2 and 3 will be neighbors, but they're invalid until `expand_indent`
+    # is called on them.
+    #
+    # When we are adding code within a method or class (at the same indentation level),
+    # use the empty lines to denote the programmer intended logical chunks.
+    # Stop and check each one. For example:
+    #
+    #   1 def dog
+    #   2   print "dog"
+    #   3
+    #   4   hash = {
+    #   5 end
+    #
+    # If we did not stop parsing at empty newlines then the block might mistakenly grab all
+    # the contents (lines 2, 3, and 4) and report them as being problems, instead of only
+    # line 4.
+    #
+    # ## Case #2: Expand/grab other logical blocks
+    #
+    # Once the search algorithm has converted all lines into blocks at a given indentation
+    # it will then `expand_indent`. Once the blocks that generates are expanded as neighbors
+    # we then begin seeing neighbors being other logical blocks i.e. a block's neighbors
+    # may be another method or class (something with keywords/ends).
+    #
+    # For example:
+    #
+    #   1 def bark
+    #   2
+    #   3 end
+    #   4
+    #   5 def sit
+    #   6 end
+    #
+    # In this case if lines 4, 5, and 6 are in a block when it tries to expand neighbors
+    # it will expand up. If it stops after line 2 or 3 it may cause problems since there's a
+    # valid kw/end pair, but the block will be checked without it.
+    #
+    # We try to resolve this edge case with `lookahead_balance_one_line` below.
     def expand_neighbors(block)
-      expanded_lines = AroundBlockScan.new(code_lines: @code_lines, block: block)
+      neighbors = AroundBlockScan.new(code_lines: @code_lines, block: block)
         .skip(:hidden?)
         .stop_after_kw
-        .scan_neighbors
-        .scan_while { |line| line.empty? } # Slurp up empties
+        .scan_neighbors_not_empty
+
+      # Slurp up empties
+      with_empties = neighbors
+        .scan_while { |line| line.empty? }
+
+      # If next line is kw and it will balance us, take it
+      expanded_lines = with_empties
+        .lookahead_balance_one_line
         .lines
 
+      # Don't allocate a block if it won't be used
+      #
+      # If nothing was taken, return nil to indicate that status
+      # used in `def call` to determine if
+      # we need to expand up/out (`expand_indent`)
       if block.lines == expanded_lines
         nil
       else
diff --git a/lib/syntax_suggest/capture_code_context.rb b/lib/syntax_suggest/capture_code_context.rb
index 7d6a550612..547072e2bd 100644
--- a/lib/syntax_suggest/capture_code_context.rb
+++ b/lib/syntax_suggest/capture_code_context.rb
@@ -76,7 +76,6 @@ module SyntaxSuggest
     #     end
     #   end
     #
-    #
     def capture_falling_indent(block)
       AroundBlockScan.new(
         block: block,