1 files changed, 101 insertions, 14 deletions

diff --git a/lib/prism/parse_result/newlines.rb b/lib/prism/parse_result/newlines.rb
index 03acb0b862..cc1343dfda 100644
--- a/lib/prism/parse_result/newlines.rb
+++ b/lib/prism/parse_result/newlines.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Prism
-  class ParseResult
+  class ParseResult < Result
     # The :line tracepoint event gets fired whenever the Ruby VM encounters an
     # expression on a new line. The types of expressions that can trigger this
     # event are:
@@ -17,21 +17,26 @@ module Prism
     # Note that the logic in this file should be kept in sync with the Java
     # MarkNewlinesVisitor, since that visitor is responsible for marking the
     # newlines for JRuby/TruffleRuby.
+    #
+    # This file is autoloaded only when `mark_newlines!` is called, so the
+    # re-opening of the various nodes in this file will only be performed in
+    # that case. We do that to avoid storing the extra `@newline` instance
+    # variable on every node if we don't need it.
     class Newlines < Visitor
       # Create a new Newlines visitor with the given newline offsets.
-      def initialize(newline_marked)
-        @newline_marked = newline_marked
+      def initialize(lines)
+        @lines = Array.new(1 + lines, false)
       end
 
       # Permit block/lambda nodes to mark newlines within themselves.
       def visit_block_node(node)
-        old_newline_marked = @newline_marked
-        @newline_marked = Array.new(old_newline_marked.size, false)
+        old_lines = @lines
+        @lines = Array.new(old_lines.size, false)
 
         begin
           super(node)
         ensure
-          @newline_marked = old_newline_marked
+          @lines = old_lines
         end
       end
 
@@ -39,7 +44,7 @@ module Prism
 
       # Mark if/unless nodes as newlines.
       def visit_if_node(node)
-        node.set_newline_flag(@newline_marked)
+        node.newline!(@lines)
         super(node)
       end
 
@@ -48,19 +53,101 @@ module Prism
       # Permit statements lists to mark newlines within themselves.
       def visit_statements_node(node)
         node.body.each do |child|
-          child.set_newline_flag(@newline_marked)
+          child.newline!(@lines)
         end
         super(node)
       end
     end
+  end
+
+  class Node
+    def newline? # :nodoc:
+      @newline ? true : false
+    end
+
+    def newline!(lines) # :nodoc:
+      line = location.start_line
+      unless lines[line]
+        lines[line] = true
+        @newline = true
+      end
+    end
+  end
+
+  class BeginNode < Node
+    def newline!(lines) # :nodoc:
+      # Never mark BeginNode with a newline flag, mark children instead.
+    end
+  end
+
+  class ParenthesesNode < Node
+    def newline!(lines) # :nodoc:
+      # Never mark ParenthesesNode with a newline flag, mark children instead.
+    end
+  end
+
+  class IfNode < Node
+    def newline!(lines) # :nodoc:
+      predicate.newline!(lines)
+    end
+  end
+
+  class UnlessNode < Node
+    def newline!(lines) # :nodoc:
+      predicate.newline!(lines)
+    end
+  end
+
+  class UntilNode < Node
+    def newline!(lines) # :nodoc:
+      predicate.newline!(lines)
+    end
+  end
+
+  class WhileNode < Node
+    def newline!(lines) # :nodoc:
+      predicate.newline!(lines)
+    end
+  end
+
+  class RescueModifierNode < Node
+    def newline!(lines) # :nodoc:
+      expression.newline!(lines)
+    end
+  end
+
+  class InterpolatedMatchLastLineNode < Node
+    def newline!(lines) # :nodoc:
+      first = parts.first
+      first.newline!(lines) if first
+    end
+  end
+
+  class InterpolatedRegularExpressionNode < Node
+    def newline!(lines) # :nodoc:
+      first = parts.first
+      first.newline!(lines) if first
+    end
+  end
+
+  class InterpolatedStringNode < Node
+    def newline!(lines) # :nodoc:
+      first = parts.first
+      first.newline!(lines) if first
+    end
+  end
 
-    private_constant :Newlines
+  class InterpolatedSymbolNode < Node
+    def newline!(lines) # :nodoc:
+      first = parts.first
+      first.newline!(lines) if first
+    end
+  end
 
-    # Walk the tree and mark nodes that are on a new line.
-    def mark_newlines!
-      value = self.value
-      raise "This method should only be called on a parse result that contains a node" unless Node === value
-      value.accept(Newlines.new(Array.new(1 + source.offsets.size, false))) # steep:ignore
+  class InterpolatedXStringNode < Node
+    def newline!(lines) # :nodoc:
+      first = parts.first
+      first.newline!(lines) if first
     end
   end
 end