2 files changed, 76 insertions, 48 deletions
diff --git a/lib/error_highlight/formatter.rb b/lib/error_highlight/formatter.rb
index b4fce93e55..4d5c82f460 100644
--- a/lib/error_highlight/formatter.rb
+++ b/lib/error_highlight/formatter.rb
@@ -2,55 +2,46 @@ module ErrorHighlight
   class DefaultFormatter
     def self.message_for(spot)
       # currently only a one-line code snippet is supported
-      if spot[:first_lineno] == spot[:last_lineno]
-        spot = truncate(spot)
+      return "" unless spot[:first_lineno] == spot[:last_lineno]
 
-        indent = spot[:snippet][0...spot[:first_column]].gsub(/[^\t]/, " ")
-        marker = indent + "^" * (spot[:last_column] - spot[:first_column])
+      snippet      = spot[:snippet]
+      first_column = spot[:first_column]
+      last_column  = spot[:last_column]
 
-        "\n\n#{ spot[:snippet] }#{ marker }"
-      else
-        ""
+      # truncate snippet to fit in the viewport
+      if snippet.size > viewport_size
+        visible_start = [first_column - viewport_size / 2, 0].max
+        visible_end   = visible_start + viewport_size
+
+        # avoid centering the snippet when the error is at the end of the line
+        visible_start = snippet.size - viewport_size if visible_end > snippet.size
+
+        prefix = visible_start.positive?    ? "..." : ""
+        suffix = visible_end < snippet.size ? "..." : ""
+
+        snippet = prefix + snippet[(visible_start + prefix.size)...(visible_end - suffix.size)] + suffix
+        snippet << "\n" unless snippet.end_with?("\n")
+
+        first_column = first_column - visible_start
+        last_column  = [last_column - visible_start, snippet.size - 1].min
       end
+
+      indent = snippet[0...first_column].gsub(/[^\t]/, " ")
+      marker = indent + "^" * (last_column - first_column)
+
+      "\n\n#{ snippet }#{ marker }"
     end
 
     def self.viewport_size
-      Ractor.current[:__error_highlight_viewport_size__] || terminal_columns
+      Ractor.current[:__error_highlight_viewport_size__] ||= terminal_columns
     end
 
     def self.viewport_size=(viewport_size)
       Ractor.current[:__error_highlight_viewport_size__] = viewport_size
     end
 
-    private
-    
-    def self.truncate(spot)
-      ellipsis = '...'
-      snippet = spot[:snippet]
-      diff = snippet.size - (viewport_size - ellipsis.size)
-
-      # snippet fits in the terminal
-      return spot if diff.negative?
-
-      if spot[:first_column] < diff
-        snippet = snippet[0...snippet.size - diff]
-        {
-          **spot,
-          snippet: snippet + ellipsis + "\n",
-          last_column: [spot[:last_column], snippet.size].min
-        }
-      else
-        {
-          **spot,
-          snippet: ellipsis + snippet[diff..-1],
-          first_column: spot[:first_column] - (diff - ellipsis.size),
-          last_column: spot[:last_column] - (diff - ellipsis.size)
-        }
-      end
-    end
-    
     def self.terminal_columns
-      # lazy load io/console in case viewport_size is set
+      # lazy load io/console, so it's not loaded when viewport_size is set
       require "io/console"
       IO.console.winsize[1]
     end
diff --git a/test/error_highlight/test_error_highlight.rb b/test/error_highlight/test_error_highlight.rb
index fd5e32a549..d00c04995d 100644
--- a/test/error_highlight/test_error_highlight.rb
+++ b/test/error_highlight/test_error_highlight.rb
@@ -5,6 +5,8 @@ require "did_you_mean"
 require "tempfile"
 
 class ErrorHighlightTest < Test::Unit::TestCase
+  ErrorHighlight::DefaultFormatter.viewport_size = 80
+
   class DummyFormatter
     def self.message_for(corrections)
       ""
@@ -12,8 +14,6 @@ class ErrorHighlightTest < Test::Unit::TestCase
   end
 
   def setup
-    ErrorHighlight::DefaultFormatter.viewport_size = 80
-
     if defined?(DidYouMean)
       @did_you_mean_old_formatter = DidYouMean.formatter
       DidYouMean.formatter = DummyFormatter
@@ -1287,27 +1287,64 @@ undefined method `time' for #{ ONE_RECV_MESSAGE }
     end
   end
 
-  def test_errors_on_small_viewports_when_error_lives_at_the_end
+  def test_errors_on_small_viewports_at_the_end
+    assert_error_message(NoMethodError, <<~END) do
+undefined method `time' for #{ ONE_RECV_MESSAGE }
+
+...0000000000000000000000000000000000000000000000000000000000000000 + 1.time {}
+                                                                       ^^^^^
+    END
+
+    100000000000000000000000000000000000000000000000000000000000000000000000000000 + 1.time {}
+    end
+  end
+
+  def test_errors_on_small_viewports_at_the_beginning
+    assert_error_message(NoMethodError, <<~END) do
+undefined method `time' for #{ ONE_RECV_MESSAGE }
+
+      1.time { 10000000000000000000000000000000000000000000000000000000000000...
+       ^^^^^
+    END
+
+      1.time { 100000000000000000000000000000000000000000000000000000000000000000000000000000 }
+
+    end
+  end
+
+  def test_errors_on_small_viewports_at_the_middle
+    assert_error_message(NoMethodError, <<~END) do
+undefined method `time' for #{ ONE_RECV_MESSAGE }
+
+...000000000000000000000000000000000 + 1.time { 10000000000000000000000000000...
+                                        ^^^^^
+    END
+
+    100000000000000000000000000000000000000 + 1.time { 100000000000000000000000000000000000000 }
+    end
+  end
+
+  def test_errors_on_small_viewports_when_larger_than_viewport
     assert_error_message(NoMethodError, <<~END) do
-undefined method 'gsuub' for an instance of String
+undefined method `timessssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss!' for #{ ONE_RECV_MESSAGE }
 
-...ooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo".gsuub(//, "")
-                                                                 ^^^^^^
+      1.timesssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss...
+       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     END
 
-      "fooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo".gsuub(//, "")
+      1.timessssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss!
     end
   end
 
-  def test_errors_on_small_viewports_when_error_lives_at_the_beginning
+  def test_errors_on_small_viewports_when_exact_size_of_viewport
     assert_error_message(NoMethodError, <<~END) do
-undefined method 'gsuub' for an instance of Integer
+undefined method `timessssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss!' for #{ ONE_RECV_MESSAGE }
 
-      1.gsuub(//, "fooooooooooooooooooooooooooooooooooooooooooooooooooooooooo...
-       ^^^^^^
+      1.timessssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss!...
+       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     END
 
-      1.gsuub(//, "fooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo")
+      1.timessssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss! * 1000
     end
   end