Engine: Fix newline handling after heredoc terminators

lib/herb/engine.rb

@@ -175,6 +175,14 @@ def self.css(value)
       end
     end
 
+    def self.comment?(code)
+      code.include?("#")
+    end
+
+    def self.heredoc?(code)
+      code.match?(/<<[~-]?\s*['"`]?\w/)
+    end
+
     protected
 
     def add_text(text)
@@ -196,8 +204,8 @@ def add_code(code)
         @src << " " << code
 
         # TODO: rework and check for Prism::InlineComment as soon as we expose the Prism Nodes in the Herb AST
-        if code.include?("#")
-          @src << "\n"
+        if self.class.comment?(code) || self.class.heredoc?(code)
+          @src << "\n" unless code[-1] == "\n"
         else
           @src << ";" unless code[-1] == "\n"
         end
@@ -267,20 +275,12 @@ def add_expression_block_end(code, escaped: false)
     end
 
     def trailing_newline(code)
-      return "\n" if comment?(code)
-      return "\n" if heredoc?(code)
+      return "\n" if self.class.comment?(code)
+      return "\n" if self.class.heredoc?(code)
 
       ""
     end
 
-    def comment?(code)
-      code.include?("#")
-    end
-
-    def heredoc?(code)
-      code.match?(/<<[~-]?\s*['"`]?\w/)
-    end
-
     def add_postamble(postamble)
       terminate_expression
       @src << postamble

lib/herb/engine/compiler.rb

@@ -545,7 +545,7 @@ def apply_trim(node, code)
 
         if at_line_start?
           lspace = extract_and_remove_lspace!
-          rspace = " \n"
+          rspace = Herb::Engine.heredoc?(code) ? "\n" : " \n"
 
           @tokens << [:code, "#{lspace}#{code}#{rspace}", current_context]
           @trim_next_whitespace = true

test/engine/engine_test.rb

@@ -177,5 +177,65 @@ class EngineTest < Minitest::Spec
 
       assert_compiled_snapshot(template)
     end
+
+    test "heredoc in code tag compiles to valid Ruby" do
+      template = <<~ERB
+        <%
+          text = <<~TEXT
+            Hello, world!
+          TEXT
+        %>
+      ERB
+
+      assert_compiled_snapshot(template)
+    end
+
+    test "heredoc in code tag inline compiles to valid Ruby" do
+      template = <<~ERB
+        <div><% text = <<~TEXT
+            Hello, world!
+          TEXT
+        %></div>
+      ERB
+
+      assert_compiled_snapshot(template)
+    end
+
+    test "heredoc with dash syntax in code tag compiles to valid Ruby" do
+      template = <<~ERB
+        <%
+          text = <<-TEXT
+            Hello, world!
+          TEXT
+        %>
+      ERB
+
+      assert_compiled_snapshot(template)
+    end
+
+    test "heredoc with quoted identifier in code tag compiles to valid Ruby" do
+      template = <<~ERB
+        <%
+          text = <<~'TEXT'
+            Hello, world!
+          TEXT
+        %>
+      ERB
+
+      assert_compiled_snapshot(template)
+    end
+
+    test "heredoc in escaped expression tag compiles to valid Ruby" do
+      template = <<~ERB
+        <%== method_call <<~GRAPHQL, variables
+          query {
+            field
+          }
+        GRAPHQL
+        %>
+      ERB
+
+      assert_compiled_snapshot(template)
+    end
   end
 end

test/engine/evaluation_test.rb

@@ -484,5 +484,18 @@ class EvaluationTest < Minitest::Spec
 
       assert_evaluated_snapshot(template, { some_condition: false }, { escape: false })
     end
+
+    test "heredoc in code tag" do
+      template = <<~ERB
+        <%
+          text = <<~TEXT
+            Hello, world!
+          TEXT
+        %>
+        <%= text %>
+      ERB
+
+      assert_evaluated_snapshot(template, {}, { escape: false })
+    end
   end
 end

test/parser/erb_test.rb

@@ -380,5 +380,26 @@ class ERBTest < Minitest::Spec
         <% end %>
       HTML
     end
+
+    test "heredoc in code tag" do
+      assert_parsed_snapshot(<<~HTML)
+        <%
+          text = <<~TEXT
+            Hello, world!
+          TEXT
+        %>
+      HTML
+    end
+
+    test "heredoc in output tag" do
+      assert_parsed_snapshot(<<~HTML)
+        <%= method_call <<~GRAPHQL, variables
+          query {
+            field
+          }
+        GRAPHQL
+        %>
+      HTML
+    end
   end
 end