From bcfad7dd3572a81fcf4d5ce8431bdfcae76bb018 Mon Sep 17 00:00:00 2001
From: schneems <richard.schneeman+foo@gmail.com>
Date: Fri, 4 Dec 2020 14:13:45 -0600
Subject: [PATCH 1/4] Fix expansion to cover already hidden lines

---
 lib/syntax_search/block_expand.rb | 1 +
 spec/unit/code_search_spec.rb     | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/lib/syntax_search/block_expand.rb b/lib/syntax_search/block_expand.rb
index 3950539..947b039 100644
--- a/lib/syntax_search/block_expand.rb
+++ b/lib/syntax_search/block_expand.rb
@@ -44,6 +44,7 @@ def call(block)
 
     def expand_indent(block)
       block = AroundBlockScan.new(code_lines: @code_lines, block: block)
+        .skip(:hidden?)
         .stop_after_kw
         .scan_adjacent_indent
         .code_block
diff --git a/spec/unit/code_search_spec.rb b/spec/unit/code_search_spec.rb
index fa3b85d..aa283aa 100644
--- a/spec/unit/code_search_spec.rb
+++ b/spec/unit/code_search_spec.rb
@@ -158,6 +158,8 @@ def hello
              1  require 'rails_helper'
              2
              3  RSpec.describe AclassNameHere, type: :worker do
+          ❯  4    describe "thing" do
+          ❯ 16    end # line 16 accidental end, but valid block
           ❯ 30    end # mismatched due to 16
             31  end
         EOM

From e148b542321c2f9ba51ad971b3140e9e216e27a0 Mon Sep 17 00:00:00 2001
From: schneems <richard.schneeman+foo@gmail.com>
Date: Fri, 4 Dec 2020 14:17:07 -0600
Subject: [PATCH 2/4] [close #20] Simplify large file output

This change reduces the output to only the minimal context needed.
---
 lib/syntax_search/around_block_scan.rb        | 21 ++++-
 lib/syntax_search/capture_code_context.rb     | 58 ++++++++++++
 lib/syntax_search/code_block.rb               |  4 +
 lib/syntax_search/code_line.rb                |  6 +-
 lib/syntax_search/code_search.rb              | 15 +++-
 .../display_code_with_line_numbers.rb         | 56 ++++++++++++
 lib/syntax_search/display_invalid_blocks.rb   | 54 +++++------
 spec/unit/code_search_spec.rb                 |  1 -
 spec/unit/display_invalid_blocks_spec.rb      | 89 ++++++++++++++++++-
 spec/unit/exe_spec.rb                         |  5 +-
 10 files changed, 268 insertions(+), 41 deletions(-)
 create mode 100644 lib/syntax_search/capture_code_context.rb
 create mode 100644 lib/syntax_search/display_code_with_line_numbers.rb

diff --git a/lib/syntax_search/around_block_scan.rb b/lib/syntax_search/around_block_scan.rb
index 837dbc1..f520f04 100644
--- a/lib/syntax_search/around_block_scan.rb
+++ b/lib/syntax_search/around_block_scan.rb
@@ -85,6 +85,26 @@ def scan_while(&block)
       self
     end
 
+    def on_falling_indent
+      last_indent = @orig_indent
+      before_lines.reverse.each do |line|
+        next if line.empty?
+        if line.indent < last_indent
+          yield line
+          last_indent = line.indent
+        end
+      end
+
+      last_indent = @orig_indent
+      after_lines.each do |line|
+        next if line.empty?
+        if line.indent < last_indent
+          yield line
+          last_indent = line.indent
+        end
+      end
+    end
+
     def scan_neighbors
       self.scan_while {|line| line.not_empty? && line.indent >= @orig_indent }
     end
@@ -93,7 +113,6 @@ def scan_adjacent_indent
       before_indent = @code_lines[@orig_before_index.pred]&.indent || 0
       after_indent = @code_lines[@orig_after_index.next]&.indent || 0
 
-
       indent = [before_indent, after_indent].min
       self.scan_while {|line| line.not_empty? && line.indent >= indent }
 
diff --git a/lib/syntax_search/capture_code_context.rb b/lib/syntax_search/capture_code_context.rb
new file mode 100644
index 0000000..b71e84e
--- /dev/null
+++ b/lib/syntax_search/capture_code_context.rb
@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module SyntaxErrorSearch
+
+  # Given a block, this method will capture surrounding
+  # code to give the user more context for the location of
+  # the problem.
+  #
+  # Return is an array of CodeLines to be rendered.
+  #
+  # Surrounding code is captured regardless of visible state
+  #
+  #   puts block.to_s # => "def bark"
+  #
+  #   context = CaptureCodeContext.new(
+  #     blocks: block,
+  #     code_lines: code_lines
+  #   )
+  #
+  #   puts context.call.join
+  #   # =>
+  #     class Dog
+  #       def bark
+  #     end
+  #
+  class CaptureCodeContext
+    attr_reader :code_lines
+
+    def initialize(blocks: , code_lines:)
+      @blocks = Array(blocks)
+      @code_lines = code_lines
+      @visible_lines = @blocks.map(&:visible_lines).flatten
+      @lines_to_output = @visible_lines.dup
+    end
+
+    def call
+      @blocks.each do |block|
+        around_lines = AroundBlockScan.new(code_lines: @code_lines, block: block)
+          .stop_after_kw
+          .scan_while {|line| line.empty? || line.indent >= block.current_indent }
+          .code_block
+          .lines
+          .select {|l| l.indent == block.current_indent && ! block.lines.include?(l) }
+
+        @lines_to_output.concat(around_lines)
+
+        AroundBlockScan.new(
+          block: block,
+          code_lines: @code_lines,
+        ).on_falling_indent do |line|
+          @lines_to_output << line
+        end
+      end
+
+      return @lines_to_output.uniq.select(&:not_empty?).sort
+    end
+  end
+end
diff --git a/lib/syntax_search/code_block.rb b/lib/syntax_search/code_block.rb
index d2f70aa..32d8e6b 100644
--- a/lib/syntax_search/code_block.rb
+++ b/lib/syntax_search/code_block.rb
@@ -23,6 +23,10 @@ def initialize(lines: [])
       @lines = Array(lines)
     end
 
+    def visible_lines
+      @lines.select(&:visible?).select(&:not_empty?)
+    end
+
     def mark_invisible
       @lines.map(&:mark_invisible)
     end
diff --git a/lib/syntax_search/code_line.rb b/lib/syntax_search/code_line.rb
index 4b6cce4..fc12851 100644
--- a/lib/syntax_search/code_line.rb
+++ b/lib/syntax_search/code_line.rb
@@ -29,7 +29,7 @@ module SyntaxErrorSearch
   # Marking a line as invisible also lets the overall program know
   # that it should not check that area for syntax errors.
   class CodeLine
-    attr_reader :line, :index, :indent
+    attr_reader :line, :index, :indent, :original_line
 
     def initialize(line: , index:)
       @original_line = line.freeze
@@ -60,6 +60,10 @@ def initialize(line: , index:)
       @is_end = (@end_count - @kw_count) > 0
     end
 
+    def <=>(b)
+      self.index <=> b.index
+    end
+
     def is_comment?
       @is_comment
     end
diff --git a/lib/syntax_search/code_search.rb b/lib/syntax_search/code_search.rb
index 4f01640..6370230 100644
--- a/lib/syntax_search/code_search.rb
+++ b/lib/syntax_search/code_search.rb
@@ -79,6 +79,14 @@ def push(block, name: )
       end
     end
 
+    # Removes the block without putting it back in the frontier
+    def sweep(block:, name: )
+      record(block: block, name: name)
+
+      block.lines.each(&:mark_invisible)
+      frontier.register_indent_block(block)
+    end
+
     # Parses the most indented lines into blocks that are marked
     # and added to the frontier
     def add_invalid_blocks
@@ -118,9 +126,10 @@ def sweep_heredocs
     end
 
     def sweep_comments
-      @code_lines.select(&:is_comment?).each do |line|
-        line.mark_invisible
-      end
+      lines = @code_lines.select(&:is_comment?)
+      return if lines.empty?
+      block = CodeBlock.new(lines: lines)
+      sweep(block: block, name: "comments")
     end
 
     # Main search loop
diff --git a/lib/syntax_search/display_code_with_line_numbers.rb b/lib/syntax_search/display_code_with_line_numbers.rb
new file mode 100644
index 0000000..998c57f
--- /dev/null
+++ b/lib/syntax_search/display_code_with_line_numbers.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module SyntaxErrorSearch
+  # Outputs code with highlighted lines
+  #
+  # Whatever is passed to this class will be rendered
+  # even if it is "marked invisible" any filtering of
+  # output should be done before calling this class.
+  #
+  #
+  #   DisplayCodeWithLineNumbers.new(
+  #     lines: lines,
+  #     highlight_lines: [lines[2], lines[3]]
+  #   ).call
+  #   # =>
+  #       1
+  #       2  def cat
+  #     ❯ 3    Dir.chdir
+  #     ❯ 4    end
+  #       5  end
+  #       6
+  class DisplayCodeWithLineNumbers
+    TERMINAL_HIGHLIGHT = "\e[1;3m" # Bold, italics
+    TERMINAL_END = "\e[0m"
+
+    def initialize(lines: , highlight_lines: [], terminal: false)
+      @lines = lines.sort
+      @terminal = terminal
+      @highlight_line_hash = highlight_lines.each_with_object({}) {|line, h| h[line] = true  }
+      @digit_count = @lines.last&.line_number.to_s.length
+    end
+
+    def call
+      @lines.map do |line|
+        string = String.new("")
+        if @highlight_line_hash[line]
+          string << "❯ "
+        else
+          string << "  "
+        end
+
+        number = line.line_number.to_s.rjust(@digit_count)
+        string << number.to_s
+        if line.empty?
+          string << line.original_line
+        else
+          string << "  "
+          string << TERMINAL_HIGHLIGHT if @terminal && @highlight_line_hash[line] # Bold, italics
+          string << line.original_line
+          string << TERMINAL_END if @terminal
+        end
+        string
+      end.join
+    end
+  end
+end
diff --git a/lib/syntax_search/display_invalid_blocks.rb b/lib/syntax_search/display_invalid_blocks.rb
index 95e04aa..e1f6766 100644
--- a/lib/syntax_search/display_invalid_blocks.rb
+++ b/lib/syntax_search/display_invalid_blocks.rb
@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require_relative "capture_code_context"
+require_relative "display_code_with_line_numbers"
+
 module SyntaxErrorSearch
   # Used for formatting invalid blocks
   class DisplayInvalidBlocks
@@ -11,11 +14,10 @@ def initialize(code_lines: ,blocks:, io: $stderr, filename: nil, terminal: false
       @io = io
 
       @blocks = Array(blocks)
-      @lines = @blocks.map(&:lines).flatten
+
+      @invalid_lines = @blocks.map(&:lines).flatten
       @code_lines = code_lines
-      @digit_count = @code_lines.last&.line_number.to_s.length
 
-      @invalid_line_hash = @lines.each_with_object({}) {|line, h| h[line] = true  }
       @invalid_type = invalid_type
     end
 
@@ -70,41 +72,29 @@ def indent(string, with: "    ")
 
     def code_block
       string = String.new("")
-      string << code_with_lines
+      string << code_with_context
       string
     end
 
-    def terminal_end
-      "\e[0m"
-    end
-
-    def terminal_highlight
-      "\e[1;3m" # Bold, italics
+    def code_with_context
+      lines = CaptureCodeContext.new(
+        blocks: @blocks,
+        code_lines: @code_lines
+      ).call
+
+      DisplayCodeWithLineNumbers.new(
+        lines: lines,
+        terminal: @terminal,
+        highlight_lines: @invalid_lines,
+      ).call
     end
 
     def code_with_lines
-      @code_lines.map do |line|
-        next if line.hidden?
-
-        string = String.new("")
-        if @invalid_line_hash[line]
-          string << "❯ "
-        else
-          string << "  "
-        end
-
-        number = line.line_number.to_s.rjust(@digit_count)
-        string << number.to_s
-        if line.empty?
-          string << line.to_s
-        else
-          string << "  "
-          string << terminal_highlight if @terminal && @invalid_line_hash[line] # Bold, italics
-          string << line.to_s
-          string << terminal_end if @terminal
-        end
-        string
-      end.join
+      DisplayCodeWithLineNumbers.new(
+        lines: @code_lines.select(&:visible?),
+        terminal: @terminal,
+        highlight_lines: @invalid_lines,
+      ).call
     end
   end
 end
diff --git a/spec/unit/code_search_spec.rb b/spec/unit/code_search_spec.rb
index aa283aa..9f9f6c2 100644
--- a/spec/unit/code_search_spec.rb
+++ b/spec/unit/code_search_spec.rb
@@ -166,7 +166,6 @@ def hello
       end
     end
 
-
     # For code that's not perfectly formatted, we ideally want to do our best
     # These examples represent the results that exist today, but I would like to improve upon them
     describe "needs improvement" do
diff --git a/spec/unit/display_invalid_blocks_spec.rb b/spec/unit/display_invalid_blocks_spec.rb
index ecb3be0..b472a60 100644
--- a/spec/unit/display_invalid_blocks_spec.rb
+++ b/spec/unit/display_invalid_blocks_spec.rb
@@ -4,6 +4,90 @@
 
 module SyntaxErrorSearch
   RSpec.describe DisplayInvalidBlocks do
+    it "captures surrounding context on same indent" do
+      syntax_string = <<~EOM
+        class Blerg
+        end
+        class OH
+
+          def lol
+          end
+
+          it "foo"
+          end
+
+          def haha
+          end
+
+        end
+
+        class Zerg
+        end
+      EOM
+
+      search = CodeSearch.new(syntax_string)
+      search.call
+
+      code_context = CaptureCodeContext.new(
+        blocks: search.invalid_blocks,
+        code_lines: search.code_lines
+      )
+
+      # Finds lines previously hidden
+      lines = code_context.call
+      expect(lines.select(&:hidden?).map(&:line_number)).to eq([11, 12])
+
+      out = DisplayCodeWithLineNumbers.new(
+        lines: lines,
+      ).call
+
+      expect(out).to eq(<<~EOM.indent(2))
+         3  class OH
+         5    def lol
+         6    end
+         8    it "foo"
+         9    end
+        11    def haha
+        12    end
+        14  end
+      EOM
+    end
+
+    it "captures surrounding context on falling indent" do
+      syntax_string = <<~EOM
+        class Blerg
+        end
+
+        class OH
+
+          def hello
+            it "foo" do
+          end
+        end
+
+        class Zerg
+        end
+      EOM
+
+      search = CodeSearch.new(syntax_string)
+      search.call
+
+      expect(search.invalid_blocks.join.strip).to eq('it "foo" do')
+
+      display = CaptureCodeContext.new(
+        blocks: search.invalid_blocks,
+        code_lines: search.code_lines
+      )
+      lines = display.call.sort
+      expect(lines.join).to eq(<<~EOM)
+        class OH
+          def hello
+            it "foo" do
+          end
+        end
+      EOM
+    end
+
     it "works with valid code" do
       syntax_string = <<~EOM
         class OH
@@ -72,6 +156,7 @@ def hai
          5  end
       EOM
     end
+
     it "shows terminal characters" do
       code_lines = code_line_array(<<~EOM)
         class OH
@@ -109,12 +194,12 @@ def hai
       expect(display.code_with_lines).to eq(
         [
           "  1  class OH",
-         ["❯ 2  ", display.terminal_highlight, "  def hello"].join,
+         ["❯ 2  ", DisplayCodeWithLineNumbers::TERMINAL_HIGHLIGHT, "  def hello"].join,
           "  3    def hai",
           "  4    end",
           "  5  end",
           ""
-        ].join($/ + display.terminal_end)
+        ].join($/ + DisplayCodeWithLineNumbers::TERMINAL_END)
       )
     end
   end
diff --git a/spec/unit/exe_spec.rb b/spec/unit/exe_spec.rb
index 92be5fe..0b5f54d 100644
--- a/spec/unit/exe_spec.rb
+++ b/spec/unit/exe_spec.rb
@@ -34,9 +34,12 @@ def exe(cmd)
         Pathname(file.path).write(lines.join)
 
         out = exe("#{file.path} --no-terminal")
-        expect(out.strip).to include(<<~EOM.indent(4))
+
+        expect(out).to include(<<~EOM.indent(4))
              77    class Lookups
           ❯  78      def input_modes
+            148    end
+            551  end
         EOM
       end
     end

From d4ef283cef4a4fec7c0627880322897f67051136 Mon Sep 17 00:00:00 2001
From: schneems <richard.schneeman+foo@gmail.com>
Date: Mon, 7 Dec 2020 10:23:01 -0600
Subject: [PATCH 3/4] Stable sort CodeBlocks

Stable sort means that when elements are entered into an array that have the same value, they will be sorted in the same order. For example:

```
MyValue = Struct.new(:value, :name)

[MyValue.new(1, :a), MyValue.new(1, :b), MyValue.new(1, :c)]
  .sort_by {|x| x.value }
  .map(&:name)
```

You would expect this to always produce `[:a, :b, :c]` since that's the original order that they're entered in, however since ruby's sort is not "stable" this is not the case.

This non-determinism causes tests to produce different results on different systems which leads to CI failures. We can mitigate this by sorting by the value we want first, then by some consistent secondary value. In this case I'm using line number for the first code block.
---
 lib/syntax_search/code_block.rb | 6 +++++-
 spec/unit/exe_spec.rb           | 4 +++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/lib/syntax_search/code_block.rb b/lib/syntax_search/code_block.rb
index 32d8e6b..754e952 100644
--- a/lib/syntax_search/code_block.rb
+++ b/lib/syntax_search/code_block.rb
@@ -52,7 +52,11 @@ def ends_at
     # populate an array with multiple code blocks then call `sort!`
     # on it without having to specify the sorting criteria
     def <=>(other)
-      self.current_indent <=> other.current_indent
+      out = self.current_indent <=> other.current_indent
+      return out if out != 0
+
+      # Stable sort
+      self.starts_at <=> other.starts_at
     end
 
     def current_indent
diff --git a/spec/unit/exe_spec.rb b/spec/unit/exe_spec.rb
index 0b5f54d..0863b8b 100644
--- a/spec/unit/exe_spec.rb
+++ b/spec/unit/exe_spec.rb
@@ -9,7 +9,9 @@ def exe_path
     end
 
     def exe(cmd)
-      run!("#{exe_path} #{cmd}")
+      out = run!("#{exe_path} #{cmd}")
+      puts out if ENV["DEBUG"]
+      out
     end
 
     it "parses valid code" do

From 9036905158e1ff7ce30ccb928f7325eb127eb787 Mon Sep 17 00:00:00 2001
From: schneems <richard.schneeman+foo@gmail.com>
Date: Mon, 7 Dec 2020 10:43:25 -0600
Subject: [PATCH 4/4] Fix code context capturing

Previously when we were trying to capture the "neighbors" context, we were using the same logic as when trying to "expand" a block. This logic looks for unmatched keyword pairs, so when it looks up it will stop at `def foo` but not at `def foo; end`. For grabbing context we want to stop after the first contextual pair so we cannot use that same method.

Instead a new method that is specialized to finding context is introduced to the AroundBlockScan class.

 When trying to capture code context, we don't want to look for a dangling keyword, we want the first matched pair
---
 CHANGELOG.md                              |  1 +
 lib/syntax_search/around_block_scan.rb    | 55 ++++++++++++++++++++++-
 lib/syntax_search/capture_code_context.rb | 16 ++++---
 lib/syntax_search/code_line.rb            |  6 +++
 spec/unit/display_invalid_blocks_spec.rb  | 20 ++++++---
 spec/unit/exe_spec.rb                     |  9 +++-
 6 files changed, 90 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ac08ded..a63a4df 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## HEAD (unreleased)
 
+- Simplify large file output so minimal context around the invalid section is shown (https://github.com/zombocom/syntax_search/pull/26)
 - Block expansion is now lexically aware of keywords (def/do/end etc.) (https://github.com/zombocom/syntax_search/pull/24)
 - Fix bug where not all of a source is lexed which is used in heredoc detection/removal (https://github.com/zombocom/syntax_search/pull/23)
 
diff --git a/lib/syntax_search/around_block_scan.rb b/lib/syntax_search/around_block_scan.rb
index f520f04..cc7af16 100644
--- a/lib/syntax_search/around_block_scan.rb
+++ b/lib/syntax_search/around_block_scan.rb
@@ -85,6 +85,50 @@ def scan_while(&block)
       self
     end
 
+    def capture_neighbor_context
+      lines = []
+      kw_count = 0
+      end_count = 0
+      before_lines.reverse.each do |line|
+        next if line.empty?
+        break if line.indent < @orig_indent
+        next if line.indent != @orig_indent
+
+        kw_count += 1 if line.is_kw?
+        end_count += 1 if line.is_end?
+        if kw_count != 0 && kw_count == end_count
+          lines << line
+          break
+        end
+
+        lines << line
+      end
+
+      lines.reverse!
+
+      kw_count = 0
+      end_count = 0
+      after_lines.each do |line|
+        # puts "line: #{line.number} #{line.original_line}, indent: #{line.indent}, #{line.empty?} #{line.indent == @orig_indent}"
+
+        next if line.empty?
+        break if line.indent < @orig_indent
+        next if line.indent != @orig_indent
+
+        kw_count += 1 if line.is_kw?
+        end_count += 1 if line.is_end?
+        if kw_count != 0 && kw_count == end_count
+          lines << line
+          break
+        end
+
+        lines << line
+      end
+      lines.select! {|line| !line.is_comment? }
+
+      lines
+    end
+
     def on_falling_indent
       last_indent = @orig_indent
       before_lines.reverse.each do |line|
@@ -119,16 +163,23 @@ def scan_adjacent_indent
       self
     end
 
+    def start_at_next_line
+      before_index; after_index
+      @before_index -= 1
+      @after_index += 1
+      self
+    end
+
     def code_block
       CodeBlock.new(lines: @code_lines[before_index..after_index])
     end
 
     def before_index
-      @before_index || @orig_before_index
+      @before_index ||= @orig_before_index
     end
 
     def after_index
-      @after_index || @orig_after_index
+      @after_index ||= @orig_after_index
     end
 
     private def before_lines
diff --git a/lib/syntax_search/capture_code_context.rb b/lib/syntax_search/capture_code_context.rb
index b71e84e..94d5ac6 100644
--- a/lib/syntax_search/capture_code_context.rb
+++ b/lib/syntax_search/capture_code_context.rb
@@ -36,11 +36,10 @@ def initialize(blocks: , code_lines:)
     def call
       @blocks.each do |block|
         around_lines = AroundBlockScan.new(code_lines: @code_lines, block: block)
-          .stop_after_kw
-          .scan_while {|line| line.empty? || line.indent >= block.current_indent }
-          .code_block
-          .lines
-          .select {|l| l.indent == block.current_indent && ! block.lines.include?(l) }
+          .start_at_next_line
+          .capture_neighbor_context
+
+        around_lines -= block.lines
 
         @lines_to_output.concat(around_lines)
 
@@ -52,7 +51,12 @@ def call
         end
       end
 
-      return @lines_to_output.uniq.select(&:not_empty?).sort
+      @lines_to_output.select!(&:not_empty?)
+      @lines_to_output.select!(&:not_comment?)
+      @lines_to_output.uniq!
+      @lines_to_output.sort!
+
+      return @lines_to_output
     end
   end
 end
diff --git a/lib/syntax_search/code_line.rb b/lib/syntax_search/code_line.rb
index fc12851..423e1e2 100644
--- a/lib/syntax_search/code_line.rb
+++ b/lib/syntax_search/code_line.rb
@@ -68,6 +68,10 @@ def is_comment?
       @is_comment
     end
 
+    def not_comment?
+      !is_comment?
+    end
+
     def is_kw?
       @is_kw
     end
@@ -107,6 +111,8 @@ def line_number
       index + 1
     end
 
+    alias :number :line_number
+
     def not_empty?
       !empty?
     end
diff --git a/spec/unit/display_invalid_blocks_spec.rb b/spec/unit/display_invalid_blocks_spec.rb
index b472a60..dbf88ee 100644
--- a/spec/unit/display_invalid_blocks_spec.rb
+++ b/spec/unit/display_invalid_blocks_spec.rb
@@ -10,15 +10,21 @@ class Blerg
         end
         class OH
 
+          def nope
+          end
+
           def lol
           end
 
           it "foo"
+            puts "here"
           end
 
           def haha
           end
 
+          def nope
+          end
         end
 
         class Zerg
@@ -35,7 +41,7 @@ class Zerg
 
       # Finds lines previously hidden
       lines = code_context.call
-      expect(lines.select(&:hidden?).map(&:line_number)).to eq([11, 12])
+      # expect(lines.select(&:hidden?).map(&:line_number)).to eq([11, 12])
 
       out = DisplayCodeWithLineNumbers.new(
         lines: lines,
@@ -43,13 +49,13 @@ class Zerg
 
       expect(out).to eq(<<~EOM.indent(2))
          3  class OH
-         5    def lol
-         6    end
-         8    it "foo"
+         8    def lol
          9    end
-        11    def haha
-        12    end
-        14  end
+        11    it "foo"
+        13    end
+        15    def haha
+        16    end
+        20  end
       EOM
     end
 
diff --git a/spec/unit/exe_spec.rb b/spec/unit/exe_spec.rb
index 0863b8b..e3b9f3c 100644
--- a/spec/unit/exe_spec.rb
+++ b/spec/unit/exe_spec.rb
@@ -38,9 +38,14 @@ def exe(cmd)
         out = exe("#{file.path} --no-terminal")
 
         expect(out).to include(<<~EOM.indent(4))
-             77    class Lookups
+             16  class Rexe
+             40    class Options < Struct.new(
+             71    end
+          ❯  77    class Lookups
           ❯  78      def input_modes
-            148    end
+          ❯ 148    end
+            152    class CommandLineParser
+            418    end
             551  end
         EOM
       end