1 files changed, 100 insertions, 0 deletions
diff --git a/test/prism/newline_test.rb b/test/prism/newline_test.rb
new file mode 100644
index 0000000000..d31fb89bc6
--- /dev/null
+++ b/test/prism/newline_test.rb
@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require_relative "test_helper"
+
+return unless defined?(RubyVM::InstructionSequence)
+
+module Prism
+  class NewlineTest < TestCase
+    base = File.expand_path("../", __FILE__)
+    filepaths = Dir["*.rb", base: base] - %w[
+      encoding_test.rb
+      errors_test.rb
+      parser_test.rb
+      regexp_test.rb
+      static_literals_test.rb
+      unescape_test.rb
+    ]
+
+    filepaths.each do |relative|
+      define_method("test_newline_flags_#{relative}") do
+        assert_newlines(base, relative)
+      end
+    end
+
+    private
+
+    def assert_newlines(base, relative)
+      filepath = File.join(base, relative)
+      source = File.read(filepath, binmode: true, external_encoding: Encoding::UTF_8)
+      expected = rubyvm_lines(source)
+
+      result = Prism.parse_file(filepath)
+      assert_empty result.errors
+      actual = prism_lines(result)
+
+      source.each_line.with_index(1) do |line, line_number|
+        # Lines like `while (foo = bar)` result in two line flags in the
+        # bytecode but only one newline flag in the AST. We need to remove the
+        # extra line flag from the bytecode to make the test pass.
+        if line.match?(/while \(/)
+          index = expected.index(line_number)
+          expected.delete_at(index) if index
+        end
+
+        # Lines like `foo =` where the value is on the next line result in
+        # another line flag in the bytecode but only one newline flag in the
+        # AST.
+        if line.match?(/^\s+\w+ =$/)
+          if source.lines[line_number].match?(/^\s+case/)
+            actual[actual.index(line_number)] += 1
+          else
+            actual.delete_at(actual.index(line_number))
+          end
+        end
+
+        if line.match?(/^\s+\w+ = \[$/)
+          if !expected.include?(line_number) && !expected.include?(line_number + 2)
+            actual[actual.index(line_number)] += 1
+          end
+        end
+      end
+
+      assert_equal expected, actual
+    end
+
+    def ignore_warnings
+      previous_verbosity = $VERBOSE
+      $VERBOSE = nil
+      yield
+    ensure
+      $VERBOSE = previous_verbosity
+    end
+
+    def rubyvm_lines(source)
+      queue = [ignore_warnings { RubyVM::InstructionSequence.compile(source) }]
+      lines = []
+
+      while iseq = queue.shift
+        lines.concat(iseq.trace_points.filter_map { |line, event| line if event == :line })
+        iseq.each_child { |insn| queue << insn unless insn.label.start_with?("ensure in ") }
+      end
+
+      lines.sort
+    end
+
+    def prism_lines(result)
+      result.mark_newlines!
+
+      queue = [result.value]
+      newlines = []
+
+      while node = queue.shift
+        queue.concat(node.compact_child_nodes)
+        newlines << result.source.line(node.location.start_offset) if node&.newline?
+      end
+
+      newlines.sort
+    end
+  end
+end