Fix diff syntax highlighting

Refactored HTML parser to avoid duplication of newline parsing.

Closes #2235

3 files changed, 38 insertions, 21 deletions

diff --git a/lib/rouge/formatters/html_gitlab.rb b/lib/rouge/formatters/html_gitlab.rb
index 3f92212243d..6762ca47c32 100644
--- a/lib/rouge/formatters/html_gitlab.rb
+++ b/lib/rouge/formatters/html_gitlab.rb
@@ -93,16 +93,27 @@ module Rouge
       end
 
       def process_tokens(tokens)
-        num_lines = 0
-        last_val = ''
-        rendered = ''
+        rendered = []
+        current_line = ''
 
         tokens.each do |tok, val|
-          last_val = val
-          num_lines += val.scan(/\n/).size
-          rendered << span(tok, val)
+          # In the case of multi-line values (e.g. comments), we need to apply
+          # styling to each line since span elements are inline.
+          val.lines.each do |line|
+            stripped = line.chomp
+            current_line << span(tok, stripped)
+
+            if line.end_with?("\n")
+              rendered << current_line
+              current_line = ''
+            end
+          end
         end
 
+        # Add leftover text
+        rendered << current_line if current_line.present?
+
+        num_lines = rendered.size
         numbers = (@linenostart..num_lines + @linenostart - 1).to_a
 
         { numbers: numbers, code: rendered }
@@ -117,9 +128,8 @@ module Rouge
         numbers.join("\n")
       end
 
-      def wrap_lines(rendered)
+      def wrap_lines(lines)
         if @lineanchors
-          lines = rendered.split("\n")
           lines = lines.each_with_index.map do |line, index|
             number = index + @linenostart
 
@@ -136,24 +146,17 @@ module Rouge
           lines.join("\n")
         else
           if @linenos == 'inline'
-            lines = rendered.split("\n")
             lines = lines.each_with_index.map do |line, index|
               number = index + @linenostart
               "<span class=\"linenos\">#{number}</span>#{line}"
             end
             lines.join("\n")
           else
-            rendered
+            lines.join("\n")
           end
         end
       end
 
-      def wrap_values(val, element)
-        lines = val.split("\n")
-        lines = lines.map{ |x| "<span #{element}>#{x}</span>" }
-        lines.join("\n")
-      end
-
       def span(tok, val)
         # http://stackoverflow.com/a/1600584/2587286
         val = CGI.escapeHTML(val)
@@ -161,13 +164,11 @@ module Rouge
         if tok.shortname.empty?
           val
         else
-          # In the case of multi-line values (e.g. comments), we need to apply
-          # styling to each line since span elements are inline.
           if @inline_theme
             rules = @inline_theme.style_for(tok).rendered_rules
-            wrap_values(val, "style=\"#{rules.to_a.join(';')}\"")
+            "<span style=\"#{rules.to_a.join(';')}\"#{val}</span>"
           else
-            wrap_values(val, "class=\"#{tok.shortname}\"")
+            "<span class=\"#{tok.shortname}\">#{val}</span>"
           end
         end
       end
diff --git a/spec/helpers/blob_helper_spec.rb b/spec/helpers/blob_helper_spec.rb
index 76009c36099..b8bba36439a 100644
--- a/spec/helpers/blob_helper_spec.rb
+++ b/spec/helpers/blob_helper_spec.rb
@@ -47,5 +47,21 @@ describe BlobHelper do
       expect(lines[1].text).to eq('        This is line 2.')
       expect(lines[2].text).to eq('        """')
     end
+
+    context 'diff highlighting' do
+      let(:blob_name) { 'test.diff' }
+      let(:blob_content) { "+aaa\n+bbb\n- ccc\n ddd\n"}
+      let(:expected) do
+        %q(<span id="LC1" class="line"><span class="gi">+aaa</span></span>
+<span id="LC2" class="line"><span class="gi">+bbb</span></span>
+<span id="LC3" class="line"><span class="gd">- ccc</span></span>
+<span id="LC4" class="line"> ddd</span>)
+      end
+
+      it 'should highlight each line properly' do
+        result = highlight(blob_name, blob_content, nowrap: true, continue: false)
+        expect(result).to eq(expected)
+      end
+    end
   end
 end
diff --git a/spec/helpers/events_helper_spec.rb b/spec/helpers/events_helper_spec.rb
index b392371deb4..da58ab98462 100644
--- a/spec/helpers/events_helper_spec.rb
+++ b/spec/helpers/events_helper_spec.rb
@@ -58,7 +58,7 @@ describe EventsHelper do
     expected = '<pre class="code highlight white ruby">' \
       "<code><span class=\"k\">def</span> <span class=\"nf\">test</span>\n" \
       "  <span class=\"s1\">\'hello world\'</span>\n" \
-      "<span class=\"k\">end</span>\n" \
+      "<span class=\"k\">end</span>" \
       '</code></pre>'
     expect(event_note(input)).to eq(expected)
   end