Fix heredoc comma issue

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ## Fixed
 - Fix printing of extraneous trailing semicolons from inline classes, modules, methods (issue [#59](https://github.com/ruby-formatter/rufo/issues/59))
 - Fix unhandled white space between array getter/setters `x[0] [0]` (issue [#62](https://github.com/ruby-formatter/rufo/issues/62))
+- Fix comma printing for heredocs when inside a hash (issue [#61](https://github.com/ruby-formatter/rufo/issues/61))
 
 ## [0.2.0] - 2017-11-13
 ### Changed

lib/rufo/formatter.rb

@@ -159,6 +159,16 @@ def initialize(code, **options)
     # This is [[line, original_line], ...]
     @inline_declarations = []
 
+    # This is used to track how far deep we are in the AST.
+    # This is useful as it allows you to check if you are inside an array
+    # when dealing with heredocs.
+    @node_level = 0
+
+    # This represents the node level of the most recent literal elements list.
+    # It is used to track if we are in a list of elements so that commas
+    # can be added appropriately for heredocs for example.
+    @literal_elements_level = nil
+
     init_settings(options)
   end
 
@@ -175,6 +185,7 @@ def format
   end
 
   def visit(node)
+    @node_level += 1
     unless node.is_a?(Array)
       bug "unexpected node: #{node} at #{current_token}"
     end
@@ -473,6 +484,8 @@ def visit(node)
     else
       bug "Unhandled node: #{node.first}"
     end
+  ensure
+    @node_level -= 1
   end
 
   def visit_exps(exps, with_indent: false, with_lines: true, want_trailing_multiline: false)
@@ -574,7 +587,20 @@ def visit_string_literal(node)
       # Accumulate heredoc: we'll write it once
       # we find a newline.
       @heredocs << [node, tilde]
-      next_token
+      # Get the next_token while capturing any output.
+      # This is needed so that we can add a comma if one is not already present.
+      captured_output = capture_output { next_token }
+
+      inside_literal_elements_list = !@literal_elements_level.nil? &&
+                                     (@node_level - @literal_elements_level) == 2
+      needs_comma = !comma? && trailing_commas
+
+      if inside_literal_elements_list && needs_comma
+        write ','
+        @last_was_heredoc = true
+      end
+
+      @output << captured_output
       return
     elsif current_token_kind == :on_backtick
       consume_token :on_backtick
@@ -2601,6 +2627,8 @@ def visit_literal_elements(elements, inside_hash: false, inside_array: false, to
     first_space = nil
 
     elements.each_with_index do |elem, i|
+      @literal_elements_level = @node_level
+
       is_last = last?(i, elements)
       wrote_comma = false
 
@@ -2642,9 +2670,10 @@ def visit_literal_elements(elements, inside_hash: false, inside_array: false, to
         write_space unless is_last
       end
     end
+    @literal_elements_level = nil
 
     if needs_trailing_comma
-      write "," unless wrote_comma || !trailing_commas
+      write "," unless wrote_comma || !trailing_commas || @last_was_heredoc
 
       consume_end_of_line(first_space: first_space)
       write_indent
@@ -2676,7 +2705,7 @@ def visit_literal_elements(elements, inside_hash: false, inside_array: false, to
 
   def check_heredocs_in_literal_elements(is_last, needs_trailing_comma, wrote_comma)
     if (newline? || comment?) && !@heredocs.empty?
-      if !is_last || needs_trailing_comma
+      if is_last && trailing_commas
         write "," unless wrote_comma
         wrote_comma = true
       end
@@ -3339,6 +3368,15 @@ def maybe_indent(toggle, indent_size)
     end
   end
 
+  def capture_output
+    old_output = @output
+    @output = ''.dup
+    yield
+    result = @output
+    @output = old_output
+    result
+  end
+
   def write(value)
     @output << value
     @last_was_newline = false

spec/lib/rufo/formatter_source_specs/heredoc.rb.spec

@@ -192,6 +192,90 @@ EOF
   bar
 EOF
 
+#~# ORIGINAL heredoc_value_trailing_comma_already_present
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
+#~# EXPECTED
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
+#~# ORIGINAL heredoc_value_trailing_comma_already_present
+#~# trailing_commas: false
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
+#~# EXPECTED
+
+{
+  content: <<-EOF
+    This is a heredoc
+  EOF
+}
+
+#~# ORIGINAL heredoc_value_trailing_comma_not_present
+#~# trailing_commas: false
+
+{
+  content: <<-EOF
+    This is a heredoc
+  EOF
+}
+
+#~# EXPECTED
+
+{
+  content: <<-EOF
+    This is a heredoc
+  EOF
+}
+
+#~# ORIGINAL heredoc_value_trailing_comma_already_present
+#~# trailing_commas: true
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
+#~# EXPECTED
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
+#~# ORIGINAL heredoc_value_trailing_comma_not_present
+#~# trailing_commas: true
+
+{
+  content: <<-EOF
+    This is a heredoc
+  EOF
+}
+
+#~# EXPECTED
+
+{
+  content: <<-EOF,
+    This is a heredoc
+  EOF
+}
+
 #~# ORIGINAL heredoc_as_hash_value
 
 {1 => <<EOF,