Repository: ruby/prism
Pull request #2428 (2024-02-15 19:36:27)

Speed up creating Ruby AST
When creating the Ruby AST, we were previously allocating Location objects for every node and every inner location. Instead, this commit changes it to pack both the start offset and length into a single u64 and pass that into the nodes. Then, when the locations are requested via a reader method, we lazily allocate the Location objects.

This pretty drastically improves speed and memory. With this script:

```ruby
require "bundler/setup"
require "prism"

filepaths = Dir[File.expand_path("path/to/extremely/large/ruby/codebase/**/*.rb", __dir__)]
filepaths.each { |filepath| Prism.parse_file(filepath) }
puts `ps -o rss= -p #{$$}`
```

We get the following results. On `main`:

```
145152
ruby bin/bench  26.96s user 2.53s system 94% cpu 31.059 total

146016
ruby bin/bench  27.00s user 2.58s system 93% cpu 31.570 total

156048
ruby bin/bench  26.87s user 2.55s system 93% cpu 31.336 total
```

On `main` with `--yjit`:

```
166384
ruby --yjit bin/bench  17.27s user 2.55s system 91% cpu 21.622 total

164704
ruby --yjit bin/bench  17.31s user 2.53s system 92% cpu 21.495 total

154672
ruby --yjit bin/bench  17.63s user 2.72s system 91% cpu 22.302 total
```

On this branch:

```
132816
ruby bin/bench  15.06s user 2.54s system 90% cpu 19.538 total

127344
ruby bin/bench  15.05s user 2.45s system 91% cpu 19.056 total

129152
ruby bin/bench  15.13s user 2.49s system 91% cpu 19.202 total
```

On this branch with `--yjit`:

```
140688
ruby --yjit bin/bench  10.32s user 2.46s system 88% cpu 14.496 total

137888
ruby --yjit bin/bench  10.32s user 2.43s system 88% cpu 14.403 total

138704
ruby --yjit bin/bench  10.32s user 2.46s system 89% cpu 14.355 total
```


Patch:
diff --git a/lib/prism/desugar_compiler.rb b/lib/prism/desugar_compiler.rb
index daca02f3..00b3097f 100644
--- a/lib/prism/desugar_compiler.rb
+++ b/lib/prism/desugar_compiler.rb
@@ -1,6 +1,216 @@
 # frozen_string_literal: true
 
 module Prism
+  class DesugarAndWriteNode # :nodoc:
+    attr_reader :node, :source, :read_class, :write_class, :arguments
+
+    def initialize(node, source, read_class, write_class, *arguments)
+      @node = node
+      @source = source
+      @read_class = read_class
+      @write_class = write_class
+      @arguments = arguments
+    end
+
+    # Desugar `x &&= y` to `x && x = y`
+    def compile
+      AndNode.new(
+        source,
+        read_class.new(source, *arguments, node.name_loc),
+        write_class.new(source, *arguments, node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+  end
+
+  class DesugarOrWriteDefinedNode # :nodoc:
+    attr_reader :node, :source, :read_class, :write_class, :arguments
+
+    def initialize(node, source, read_class, write_class, *arguments)
+      @node = node
+      @source = source
+      @read_class = read_class
+      @write_class = write_class
+      @arguments = arguments
+    end
+
+    # Desugar `x ||= y` to `defined?(x) ? x : x = y`
+    def compile
+      IfNode.new(
+        source,
+        node.operator_loc,
+        DefinedNode.new(source, nil, read_class.new(source, *arguments, node.name_loc), nil, node.operator_loc, node.name_loc),
+        node.operator_loc,
+        StatementsNode.new(source, [read_class.new(source, *arguments, node.name_loc)], node.location),
+        ElseNode.new(
+          source,
+          node.operator_loc,
+          StatementsNode.new(
+            source,
+            [write_class.new(source, *arguments, node.name_loc, node.value, node.operator_loc, node.location)],
+            node.location
+          ),
+          node.operator_loc,
+          node.location
+        ),
+        node.operator_loc,
+        node.location
+      )
+    end
+  end
+
+  class DesugarOperatorWriteNode # :nodoc:
+    attr_reader :node, :source, :read_class, :write_class, :arguments
+
+    def initialize(node, source, read_class, write_class, *arguments)
+      @node = node
+      @source = source
+      @read_class = read_class
+      @write_class = write_class
+      @arguments = arguments
+    end
+
+    # Desugar `x += y` to `x = x + y`
+    def compile
+      write_class.new(
+        source,
+        *arguments,
+        node.name_loc,
+        CallNode.new(
+          source,
+          0,
+          read_class.new(source, *arguments, node.name_loc),
+          nil,
+          node.operator_loc.slice.chomp("="),
+          node.operator_loc.copy(length: node.operator_loc.length - 1),
+          nil,
+          ArgumentsNode.new(source, 0, [node.value], node.value.location),
+          nil,
+          nil,
+          node.location
+        ),
+        node.operator_loc.copy(start_offset: node.operator_loc.end_offset - 1, length: 1),
+        node.location
+      )
+    end
+  end
+
+  class DesugarOrWriteNode # :nodoc:
+    attr_reader :node, :source, :read_class, :write_class, :arguments
+
+    def initialize(node, source, read_class, write_class, *arguments)
+      @node = node
+      @source = source
+      @read_class = read_class
+      @write_class = write_class
+      @arguments = arguments
+    end
+
+    # Desugar `x ||= y` to `x || x = y`
+    def compile
+      OrNode.new(
+        source,
+        read_class.new(source, *arguments, node.name_loc),
+        write_class.new(source, *arguments, node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+  end
+
+  private_constant :DesugarAndWriteNode, :DesugarOrWriteNode, :DesugarOrWriteDefinedNode, :DesugarOperatorWriteNode
+
+  class ClassVariableAndWriteNode
+    def desugar # :nodoc:
+      DesugarAndWriteNode.new(self, source, ClassVariableReadNode, ClassVariableWriteNode, name).compile
+    end
+  end
+
+  class ClassVariableOrWriteNode
+    def desugar # :nodoc:
+      DesugarOrWriteDefinedNode.new(self, source, ClassVariableReadNode, ClassVariableWriteNode, name).compile
+    end
+  end
+
+  class ClassVariableOperatorWriteNode
+    def desugar # :nodoc:
+      DesugarOperatorWriteNode.new(self, source, ClassVariableReadNode, ClassVariableWriteNode, name).compile
+    end
+  end
+
+  class ConstantAndWriteNode
+    def desugar # :nodoc:
+      DesugarAndWriteNode.new(self, source, ConstantReadNode, ConstantWriteNode, name).compile
+    end
+  end
+
+  class ConstantOrWriteNode
+    def desugar # :nodoc:
+      DesugarOrWriteDefinedNode.new(self, source, ConstantReadNode, ConstantWriteNode, name).compile
+    end
+  end
+
+  class ConstantOperatorWriteNode
+    def desugar # :nodoc:
+      DesugarOperatorWriteNode.new(self, source, ConstantReadNode, ConstantWriteNode, name).compile
+    end
+  end
+
+  class GlobalVariableAndWriteNode
+    def desugar # :nodoc:
+      DesugarAndWriteNode.new(self, source, GlobalVariableReadNode, GlobalVariableWriteNode, name).compile
+    end
+  end
+
+  class GlobalVariableOrWriteNode
+    def desugar # :nodoc:
+      DesugarOrWriteDefinedNode.new(self, source, GlobalVariableReadNode, GlobalVariableWriteNode, name).compile
+    end
+  end
+
+  class GlobalVariableOperatorWriteNode
+    def desugar # :nodoc:
+      DesugarOperatorWriteNode.new(self, source, GlobalVariableReadNode, GlobalVariableWriteNode, name).compile
+    end
+  end
+
+  class InstanceVariableAndWriteNode
+    def desugar # :nodoc:
+      DesugarAndWriteNode.new(self, source, InstanceVariableReadNode, InstanceVariableWriteNode, name).compile
+    end
+  end
+
+  class InstanceVariableOrWriteNode
+    def desugar # :nodoc:
+      DesugarOrWriteNode.new(self, source, InstanceVariableReadNode, InstanceVariableWriteNode, name).compile
+    end
+  end
+
+  class InstanceVariableOperatorWriteNode
+    def desugar # :nodoc:
+      DesugarOperatorWriteNode.new(self, source, InstanceVariableReadNode, InstanceVariableWriteNode, name).compile
+    end
+  end
+
+  class LocalVariableAndWriteNode
+    def desugar # :nodoc:
+      DesugarAndWriteNode.new(self, source, LocalVariableReadNode, LocalVariableWriteNode, name, depth).compile
+    end
+  end
+
+  class LocalVariableOrWriteNode
+    def desugar # :nodoc:
+      DesugarOrWriteNode.new(self, source, LocalVariableReadNode, LocalVariableWriteNode, name, depth).compile
+    end
+  end
+
+  class LocalVariableOperatorWriteNode
+    def desugar # :nodoc:
+      DesugarOperatorWriteNode.new(self, source, LocalVariableReadNode, LocalVariableWriteNode, name, depth).compile
+    end
+  end
+
   # DesugarCompiler is a compiler that desugars Ruby code into a more primitive
   # form. This is useful for consumers that want to deal with fewer node types.
   class DesugarCompiler < MutationCompiler
@@ -10,7 +220,7 @@ module Prism
     #
     # @@foo && @@foo = bar
     def visit_class_variable_and_write_node(node)
-      desugar_and_write_node(node, ClassVariableReadNode, ClassVariableWriteNode, node.name)
+      node.desugar
     end
 
     # @@foo ||= bar
@@ -19,7 +229,7 @@ module Prism
     #
     # defined?(@@foo) ? @@foo : @@foo = bar
     def visit_class_variable_or_write_node(node)
-      desugar_or_write_defined_node(node, ClassVariableReadNode, ClassVariableWriteNode, node.name)
+      node.desugar
     end
 
     # @@foo += bar
@@ -28,7 +238,7 @@ module Prism
     #
     # @@foo = @@foo + bar
     def visit_class_variable_operator_write_node(node)
-      desugar_operator_write_node(node, ClassVariableReadNode, ClassVariableWriteNode, node.name)
+      node.desugar
     end
 
     # Foo &&= bar
@@ -37,7 +247,7 @@ module Prism
     #
     # Foo && Foo = bar
     def visit_constant_and_write_node(node)
-      desugar_and_write_node(node, ConstantReadNode, ConstantWriteNode, node.name)
+      node.desugar
     end
 
     # Foo ||= bar
@@ -46,7 +256,7 @@ module Prism
     #
     # defined?(Foo) ? Foo : Foo = bar
     def visit_constant_or_write_node(node)
-      desugar_or_write_defined_node(node, ConstantReadNode, ConstantWriteNode, node.name)
+      node.desugar
     end
 
     # Foo += bar
@@ -55,7 +265,7 @@ module Prism
     #
     # Foo = Foo + bar
     def visit_constant_operator_write_node(node)
-      desugar_operator_write_node(node, ConstantReadNode, ConstantWriteNode, node.name)
+      node.desugar
     end
 
     # $foo &&= bar
@@ -64,7 +274,7 @@ module Prism
     #
     # $foo && $foo = bar
     def visit_global_variable_and_write_node(node)
-      desugar_and_write_node(node, GlobalVariableReadNode, GlobalVariableWriteNode, node.name)
+      node.desugar
     end
 
     # $foo ||= bar
@@ -73,7 +283,7 @@ module Prism
     #
     # defined?($foo) ? $foo : $foo = bar
     def visit_global_variable_or_write_node(node)
-      desugar_or_write_defined_node(node, GlobalVariableReadNode, GlobalVariableWriteNode, node.name)
+      node.desugar
     end
 
     # $foo += bar
@@ -82,7 +292,7 @@ module Prism
     #
     # $foo = $foo + bar
     def visit_global_variable_operator_write_node(node)
-      desugar_operator_write_node(node, GlobalVariableReadNode, GlobalVariableWriteNode, node.name)
+      node.desugar
     end
 
     # @foo &&= bar
@@ -91,7 +301,7 @@ module Prism
     #
     # @foo && @foo = bar
     def visit_instance_variable_and_write_node(node)
-      desugar_and_write_node(node, InstanceVariableReadNode, InstanceVariableWriteNode, node.name)
+      node.desugar
     end
 
     # @foo ||= bar
@@ -100,7 +310,7 @@ module Prism
     #
     # @foo || @foo = bar
     def visit_instance_variable_or_write_node(node)
-      desugar_or_write_node(node, InstanceVariableReadNode, InstanceVariableWriteNode, node.name)
+      node.desugar
     end
 
     # @foo += bar
@@ -109,7 +319,7 @@ module Prism
     #
     # @foo = @foo + bar
     def visit_instance_variable_operator_write_node(node)
-      desugar_operator_write_node(node, InstanceVariableReadNode, InstanceVariableWriteNode, node.name)
+      node.desugar
     end
 
     # foo &&= bar
@@ -118,7 +328,7 @@ module Prism
     #
     # foo && foo = bar
     def visit_local_variable_and_write_node(node)
-      desugar_and_write_node(node, LocalVariableReadNode, LocalVariableWriteNode, node.name, node.depth)
+      node.desugar
     end
 
     # foo ||= bar
@@ -127,7 +337,7 @@ module Prism
     #
     # foo || foo = bar
     def visit_local_variable_or_write_node(node)
-      desugar_or_write_node(node, LocalVariableReadNode, LocalVariableWriteNode, node.name, node.depth)
+      node.desugar
     end
 
     # foo += bar
@@ -136,72 +346,7 @@ module Prism
     #
     # foo = foo + bar
     def visit_local_variable_operator_write_node(node)
-      desugar_operator_write_node(node, LocalVariableReadNode, LocalVariableWriteNode, node.name, node.depth)
-    end
-
-    private
-
-    # Desugar `x &&= y` to `x && x = y`
-    def desugar_and_write_node(node, read_class, write_class, *arguments)
-      AndNode.new(
-        read_class.new(*arguments, node.name_loc),
-        write_class.new(*arguments, node.name_loc, node.value, node.operator_loc, node.location),
-        node.operator_loc,
-        node.location
-      )
-    end
-
-    # Desugar `x += y` to `x = x + y`
-    def desugar_operator_write_node(node, read_class, write_class, *arguments)
-      write_class.new(
-        *arguments,
-        node.name_loc,
-        CallNode.new(
-          0,
-          read_class.new(*arguments, node.name_loc),
-          nil,
-          node.operator_loc.slice.chomp("="),
-          node.operator_loc.copy(length: node.operator_loc.length - 1),
-          nil,
-          ArgumentsNode.new(0, [node.value], node.value.location),
-          nil,
-          nil,
-          node.location
-        ),
-        node.operator_loc.copy(start_offset: node.operator_loc.end_offset - 1, length: 1),
-        node.location
-      )
-    end
-
-    # Desugar `x ||= y` to `x || x = y`
-    def desugar_or_write_node(node, read_class, write_class, *arguments)
-      OrNode.new(
-        read_class.new(*arguments, node.name_loc),
-        write_class.new(*arguments, node.name_loc, node.value, node.operator_loc, node.location),
-        node.operator_loc,
-        node.location
-      )
-    end
-
-    # Desugar `x ||= y` to `defined?(x) ? x : x = y`
-    def desugar_or_write_defined_node(node, read_class, write_class, *arguments)
-      IfNode.new(
-        node.operator_loc,
-        DefinedNode.new(nil, read_class.new(*arguments, node.name_loc), nil, node.operator_loc, node.name_loc),
-        node.operator_loc,
-        StatementsNode.new([read_class.new(*arguments, node.name_loc)], node.location),
-        ElseNode.new(
-          node.operator_loc,
-          StatementsNode.new(
-            [write_class.new(*arguments, node.name_loc, node.value, node.operator_loc, node.location)],
-            node.location
-          ),
-          node.operator_loc,
-          node.location
-        ),
-        node.operator_loc,
-        node.location
-      )
+      node.desugar
     end
   end
 end
diff --git a/lib/prism/parse_result.rb b/lib/prism/parse_result.rb
index bcebc1a4..5b22184e 100644
--- a/lib/prism/parse_result.rb
+++ b/lib/prism/parse_result.rb
@@ -452,17 +452,19 @@ module Prism
 
   # This represents a token from the Ruby source.
   class Token
+    # The Source object that represents the source this token came from.
+    attr_reader :source
+    private :source
+
     # The type of token that this token is.
     attr_reader :type
 
     # A byteslice of the source that this token represents.
     attr_reader :value
 
-    # A Location object representing the location of this token in the source.
-    attr_reader :location
-
     # Create a new token object with the given type, value, and location.
-    def initialize(type, value, location)
+    def initialize(source, type, value, location)
+      @source = source
       @type = type
       @value = value
       @location = location
@@ -473,6 +475,12 @@ module Prism
       { type: type, value: value, location: location }
     end
 
+    # A Location object representing the location of this token in the source.
+    def location
+      return @location if @location.is_a?(Location)
+      @location = Location.new(source, @location >> 32, @location & 0xFFFFFFFF)
+    end
+
     # Implement the pretty print interface for Token.
     def pretty_print(q)
       q.group do
diff --git a/lib/prism/translation/parser/compiler.rb b/lib/prism/translation/parser/compiler.rb
index 264d85c2..81156a64 100644
--- a/lib/prism/translation/parser/compiler.rb
+++ b/lib/prism/translation/parser/compiler.rb
@@ -1477,7 +1477,7 @@ module Prism
         # ^^^^^
         def visit_string_node(node)
           if node.opening&.start_with?("<<")
-            children, closing = visit_heredoc(InterpolatedStringNode.new(node.opening_loc, [node.copy(opening_loc: nil, closing_loc: nil, location: node.content_loc)], node.closing_loc, node.location))
+            children, closing = visit_heredoc(InterpolatedStringNode.new(node.send(:source), node.opening_loc, [node.copy(opening_loc: nil, closing_loc: nil, location: node.content_loc)], node.closing_loc, node.location))
             builder.string_compose(token(node.opening_loc), children, closing)
           elsif node.opening == "?"
             builder.character([node.unescaped, srange(node.location)])
diff --git a/src/prism.c b/src/prism.c
index 6b8473cb..e517a0cc 100644
--- a/src/prism.c
+++ b/src/prism.c
@@ -51,7 +51,6 @@ debug_context(pm_context_t context) {
         case PM_CONTEXT_IF: return "IF";
         case PM_CONTEXT_MAIN: return "MAIN";
         case PM_CONTEXT_MODULE: return "MODULE";
-        case PM_CONTEXT_NONE: return "NONE";
         case PM_CONTEXT_PARENS: return "PARENS";
         case PM_CONTEXT_POSTEXE: return "POSTEXE";
         case PM_CONTEXT_PREDICATE: return "PREDICATE";
@@ -16770,18 +16769,7 @@ parse_assignment_values(pm_parser_t *parser, pm_binding_power_t previous_binding
     if (is_single_value && match1(parser, PM_TOKEN_KEYWORD_RESCUE_MODIFIER)) {
         pm_token_t rescue = parser->current;
         parser_lex(parser);
-
-        bool accepts_command_call_inner = false;
-
-        // RHS can accept command call iff the value is a call with arguments but without paranthesis.
-        if (PM_NODE_TYPE_P(value, PM_CALL_NODE)) {
-            pm_call_node_t *call_node = (pm_call_node_t *)value;
-            if ((call_node->arguments != NULL) && (call_node->opening_loc.start == NULL)) {
-                accepts_command_call_inner = true;
-            }
-        }
-
-        pm_node_t *right = parse_expression(parser, binding_power, accepts_command_call_inner, PM_ERR_RESCUE_MODIFIER_VALUE);
+        pm_node_t *right = parse_expression(parser, binding_power, false, PM_ERR_RESCUE_MODIFIER_VALUE);
 
         return (pm_node_t *) pm_rescue_modifier_node_create(parser, value, &rescue, right);
     }
diff --git a/templates/ext/prism/api_node.c.erb b/templates/ext/prism/api_node.c.erb
index f52d537d..55867f07 100644
--- a/templates/ext/prism/api_node.c.erb
+++ b/templates/ext/prism/api_node.c.erb
@@ -12,23 +12,24 @@ static VALUE rb_cPrism<%= node.name %>;
 <%- end -%>
 
 static VALUE
-pm_location_new(pm_parser_t *parser, const uint8_t *start, const uint8_t *end, VALUE source) {
-    VALUE argv[] = { source, LONG2FIX(start - parser->start), LONG2FIX(end - start) };
-    return rb_class_new_instance(3, argv, rb_cPrismLocation);
+pm_location_new(pm_parser_t *parser, const uint8_t *start, const uint8_t *end) {
+    uint64_t value = ((((uint64_t) (start - parser->start)) << 32) | ((uint32_t) (end - start)));
+    return ULL2NUM(value);
 }
 
 VALUE
 pm_token_new(pm_parser_t *parser, pm_token_t *token, rb_encoding *encoding, VALUE source) {
     ID type = rb_intern(pm_token_type_name(token->type));
-    VALUE location = pm_location_new(parser, token->start, token->end, source);
+    VALUE location = pm_location_new(parser, token->start, token->end);
 
     VALUE argv[] = {
+        source,
         ID2SYM(type),
         rb_enc_str_new((const char *) token->start, token->end - token->start, encoding),
         location
     };
 
-    return rb_class_new_instance(3, argv, rb_cPrismToken);
+    return rb_class_new_instance(4, argv, rb_cPrismToken);
 }
 
 static VALUE
@@ -144,8 +145,11 @@ pm_ast_new(pm_parser_t *parser, pm_node_t *node, rb_encoding *encoding, VALUE so
                     <%- if node.fields.any? { |field| ![Prism::NodeField, Prism::OptionalNodeField, Prism::FlagsField].include?(field.class) } -%>
                     pm_<%= node.human %>_t *cast = (pm_<%= node.human %>_t *) node;
                     <%- end -%>
-                    VALUE argv[<%= node.fields.length + 1 %>];
-                    <%- node.fields.each_with_index do |field, index| -%>
+                    VALUE argv[<%= node.fields.length + 2 %>];
+
+                    // source
+                    argv[0] = source;
+                    <%- node.fields.each.with_index(1) do |field, index| -%>
 
                     // <%= field.name %>
                     <%- case field -%>
@@ -176,10 +180,10 @@ pm_ast_new(pm_parser_t *parser, pm_node_t *node, rb_encoding *encoding, VALUE so
                     }
                     <%- when Prism::LocationField -%>
 #line <%= __LINE__ + 1 %> "<%= File.basename(__FILE__) %>"
-                    argv[<%= index %>] = pm_location_new(parser, cast-><%= field.name %>.start, cast-><%= field.name %>.end, source);
+                    argv[<%= index %>] = pm_location_new(parser, cast-><%= field.name %>.start, cast-><%= field.name %>.end);
                     <%- when Prism::OptionalLocationField -%>
 #line <%= __LINE__ + 1 %> "<%= File.basename(__FILE__) %>"
-                    argv[<%= index %>] = cast-><%= field.name %>.start == NULL ? Qnil : pm_location_new(parser, cast-><%= field.name %>.start, cast-><%= field.name %>.end, source);
+                    argv[<%= index %>] = cast-><%= field.name %>.start == NULL ? Qnil : pm_location_new(parser, cast-><%= field.name %>.start, cast-><%= field.name %>.end);
                     <%- when Prism::UInt8Field -%>
 #line <%= __LINE__ + 1 %> "<%= File.basename(__FILE__) %>"
                     argv[<%= index %>] = UINT2NUM(cast-><%= field.name %>);
@@ -195,9 +199,9 @@ pm_ast_new(pm_parser_t *parser, pm_node_t *node, rb_encoding *encoding, VALUE so
                     <%- end -%>
 
                     // location
-                    argv[<%= node.fields.length %>] = pm_location_new(parser, node->location.start, node->location.end, source);
+                    argv[<%= node.fields.length + 1 %>] = pm_location_new(parser, node->location.start, node->location.end);
 
-                    rb_ary_push(value_stack, rb_class_new_instance(<%= node.fields.length + 1 %>, argv, rb_cPrism<%= node.name %>));
+                    rb_ary_push(value_stack, rb_class_new_instance(<%= node.fields.length + 2 %>, argv, rb_cPrism<%= node.name %>));
                     break;
                 }
                 <%- end -%>
diff --git a/templates/lib/prism/dsl.rb.erb b/templates/lib/prism/dsl.rb.erb
index be18ad45..d21cf530 100644
--- a/templates/lib/prism/dsl.rb.erb
+++ b/templates/lib/prism/dsl.rb.erb
@@ -9,10 +9,12 @@ module Prism
   #         Prism::IntegerNode.new(
   #           Prism::IntegerBaseFlags::DECIMAL,
   #           Prism::Location.new(source, 1, 1),
+  #           source
   #         )
   #       ],
   #       Prism::Location.new(source, 0, 1),
-  #       Prism::Location.new(source, 2, 1)
+  #       Prism::Location.new(source, 2, 1),
+  #       source
   #     )
   #
   # you could instead write:
@@ -20,9 +22,10 @@ module Prism
   #     source = Prism::Source.new("[1]")
   #
   #     ArrayNode(
-  #       IntegerNode(Prism::IntegerBaseFlags::DECIMAL, Location(source, 1, 1))),
+  #       IntegerNode(Prism::IntegerBaseFlags::DECIMAL, Location(source, 1, 1)), source),
   #       Location(source, 0, 1),
-  #       Location(source, 2, 1)
+  #       Location(source, 2, 1),
+  #       source
   #     )
   #
   # This is mostly helpful in the context of writing tests, but can also be used
@@ -37,8 +40,8 @@ module Prism
     <%- nodes.each do |node| -%>
 
     # Create a new <%= node.name %> node
-    def <%= node.name %>(<%= (node.fields.map(&:name) + ["location = Location()"]).join(", ") %>)
-      <%= node.name %>.new(<%= (node.fields.map(&:name) + ["location"]).join(", ") %>)
+    def <%= node.name %>(<%= (node.fields.map(&:name) + ["source = nil, location = Location()"]).join(", ") %>)
+      <%= node.name %>.new(<%= ["source", *node.fields.map(&:name), "location"].join(", ") %>)
     end
     <%- end -%>
   end
diff --git a/templates/lib/prism/node.rb.erb b/templates/lib/prism/node.rb.erb
index b74d39c0..81a007fc 100644
--- a/templates/lib/prism/node.rb.erb
+++ b/templates/lib/prism/node.rb.erb
@@ -2,9 +2,16 @@ module Prism
   # This represents a node in the tree. It is the parent class of all of the
   # various node types.
   class Node
+    # A pointer to the source that this node was created from.
+    attr_reader :source
+    private :source
+
     # A Location instance that represents the location of this node in the
     # source.
-    attr_reader :location
+    def location
+      return @location if @location.is_a?(Location)
+      @location = Location.new(source, @location >> 32, @location & 0xFFFFFFFF)
+    end
 
     def newline? # :nodoc:
       @newline ? true : false
@@ -52,7 +59,7 @@ module Prism
     # Returns an array of child nodes, including `nil`s in the place of optional
     # nodes that were not present.
     def child_nodes
-      raise NoMethodError, "undefined method `#{__method__}' for #{inspect}"
+      raise NoMethodError, "undefined method `child_nodes' for #{inspect}"
     end
 
     alias deconstruct child_nodes
@@ -81,24 +88,14 @@ module Prism
   #<%= line %>
   <%- end -%>
   class <%= node.name -%> < Node
-    <%- node.fields.each do |field| -%>
-    <%- if field.comment.nil? -%>
-    # <%= "private " if field.is_a?(Prism::FlagsField) %>attr_reader <%= field.name %>: <%= field.rbs_class %>
-    <%- else -%>
-    <%- field.each_comment_line do |line| -%>
-    #<%= line %>
-    <%- end -%>
-    <%- end -%>
-    attr_reader :<%= field.name -%><%= "\n    private :#{field.name}" if field.is_a?(Prism::FlagsField) %>
-
-    <%- end -%>
     # def initialize: (<%= (node.fields.map { |field| "#{field.rbs_class} #{field.name}" } + ["Location location"]).join(", ") %>) -> void
-    def initialize(<%= (node.fields.map(&:name) + ["location"]).join(", ") %>)
+    def initialize(source, <%= (node.fields.map(&:name) + ["location"]).join(", ") %>)
+      @source = source
       @newline = false
+      @location = location
       <%- node.fields.each do |field| -%>
       @<%= field.name %> = <%= field.name %>
       <%- end -%>
-      @location = location
     end
 
     # def accept: (Visitor visitor) -> void
@@ -173,6 +170,7 @@ module Prism
     # def copy: (**params) -> <%= node.name %>
     def copy(**params)
       <%= node.name %>.new(
+        source,
         <%- (node.fields.map(&:name) + ["location"]).map do |name| -%>
         params.fetch(:<%= name %>) { <%= name %> },
         <%- end -%>
@@ -186,6 +184,37 @@ module Prism
     def deconstruct_keys(keys)
       { <%= (node.fields.map { |field| "#{field.name}: #{field.name}" } + ["location: location"]).join(", ") %> }
     end
+
+    <%- node.fields.each do |field| -%>
+    <%- if field.comment.nil? -%>
+    # <%= "private " if field.is_a?(Prism::FlagsField) %>attr_reader <%= field.name %>: <%= field.rbs_class %>
+    <%- else -%>
+    <%- field.each_comment_line do |line| -%>
+    #<%= line %>
+    <%- end -%>
+    <%- end -%>
+    <%- case field -%>
+    <%- when Prism::LocationField -%>
+    def <%= field.name %>
+      return @<%= field.name %> if @<%= field.name %>.is_a?(Location)
+      @<%= field.name %> = Location.new(source, @<%= field.name %> >> 32, @<%= field.name %> & 0xFFFFFFFF)
+    end
+    <%- when Prism::OptionalLocationField -%>
+    def <%= field.name %>
+      case @<%= field.name %>
+      when nil
+        nil
+      when Location
+        @<%= field.name %>
+      else
+        @<%= field.name %> = Location.new(source, @<%= field.name %> >> 32, @<%= field.name %> & 0xFFFFFFFF)
+      end
+    end
+    <%- else -%>
+    attr_reader :<%= field.name -%><%= "\n    private :#{field.name}" if field.is_a?(Prism::FlagsField) %>
+    <%- end -%>
+
+    <%- end -%>
     <%- node.fields.each do |field| -%>
     <%- case field -%>
     <%- when Prism::LocationField -%>
diff --git a/templates/lib/prism/serialize.rb.erb b/templates/lib/prism/serialize.rb.erb
index 751341ce..af03a54b 100644
--- a/templates/lib/prism/serialize.rb.erb
+++ b/templates/lib/prism/serialize.rb.erb
@@ -112,7 +112,7 @@ module Prism
           length = load_varuint
           lex_state = load_varuint
           location = Location.new(@source, start, length)
-          tokens << [Prism::Token.new(type, location.slice, location), lex_state]
+          tokens << [Prism::Token.new(source, type, location.slice, location), lex_state]
         end
 
         tokens
@@ -274,7 +274,8 @@ module Prism
             <%- if node.needs_serialized_length? -%>
             load_serialized_length
             <%- end -%>
-            <%= node.name %>.new(<%= (node.fields.map { |field|
+            <%= node.name %>.new(
+              source, <%= (node.fields.map { |field|
               case field
               when Prism::NodeField then "load_node"
               when Prism::OptionalNodeField then "load_optional_node"
@@ -308,7 +309,8 @@ module Prism
               <%- if node.needs_serialized_length? -%>
               load_serialized_length
               <%- end -%>
-              <%= node.name %>.new(<%= (node.fields.map { |field|
+              <%= node.name %>.new(
+                source, <%= (node.fields.map { |field|
                 case field
                 when Prism::NodeField then "load_node"
                 when Prism::OptionalNodeField then "load_optional_node"


Test patch:
diff --git a/test/prism/fixtures/rescue.txt b/test/prism/fixtures/rescue.txt
index 7cce8663..90530b70 100644
--- a/test/prism/fixtures/rescue.txt
+++ b/test/prism/fixtures/rescue.txt
@@ -31,5 +31,3 @@ rescue
 end
 
 foo if bar rescue baz
-
-z = x y rescue c d
diff --git a/test/prism/newline_test.rb b/test/prism/newline_test.rb
index 3d42ca99..eea69ec5 100644
--- a/test/prism/newline_test.rb
+++ b/test/prism/newline_test.rb
@@ -7,7 +7,7 @@ return unless defined?(RubyVM::InstructionSequence)
 module Prism
   class NewlineTest < TestCase
     base = File.expand_path("../", __FILE__)
-    filepaths = Dir["*.rb", base: base] - %w[encoding_test.rb parser_test.rb unescape_test.rb]
+    filepaths = Dir["*.rb", base: base] - %w[encoding_test.rb errors_test.rb parser_test.rb unescape_test.rb]
 
     filepaths.each do |relative|
       define_method("test_newline_flags_#{relative}") do
diff --git a/test/prism/snapshots/rescue.txt b/test/prism/snapshots/rescue.txt
index fab9a8f7..e75fe481 100644
--- a/test/prism/snapshots/rescue.txt
+++ b/test/prism/snapshots/rescue.txt
@@ -1,8 +1,8 @@
-@ ProgramNode (location: (1,0)-(35,18))
-├── locals: [:a, :z]
+@ ProgramNode (location: (1,0)-(33,21))
+├── locals: [:a]
 └── statements:
-    @ StatementsNode (location: (1,0)-(35,18))
-    └── body: (length: 14)
+    @ StatementsNode (location: (1,0)-(33,21))
+    └── body: (length: 13)
         ├── @ RescueModifierNode (location: (1,0)-(1,14))
         │   ├── expression:
         │   │   @ CallNode (location: (1,0)-(1,3))
@@ -386,102 +386,46 @@
         │   ├── rparen_loc: ∅
         │   ├── equal_loc: ∅
         │   └── end_keyword_loc: (31,0)-(31,3) = "end"
-        ├── @ RescueModifierNode (location: (33,0)-(33,21))
-        │   ├── expression:
-        │   │   @ IfNode (location: (33,0)-(33,10))
-        │   │   ├── if_keyword_loc: (33,4)-(33,6) = "if"
-        │   │   ├── predicate:
-        │   │   │   @ CallNode (location: (33,7)-(33,10))
-        │   │   │   ├── flags: variable_call, ignore_visibility
-        │   │   │   ├── receiver: ∅
-        │   │   │   ├── call_operator_loc: ∅
-        │   │   │   ├── name: :bar
-        │   │   │   ├── message_loc: (33,7)-(33,10) = "bar"
-        │   │   │   ├── opening_loc: ∅
-        │   │   │   ├── arguments: ∅
-        │   │   │   ├── closing_loc: ∅
-        │   │   │   └── block: ∅
-        │   │   ├── then_keyword_loc: ∅
-        │   │   ├── statements:
-        │   │   │   @ StatementsNode (location: (33,0)-(33,3))
-        │   │   │   └── body: (length: 1)
-        │   │   │       └── @ CallNode (location: (33,0)-(33,3))
-        │   │   │           ├── flags: variable_call, ignore_visibility
-        │   │   │           ├── receiver: ∅
-        │   │   │           ├── call_operator_loc: ∅
-        │   │   │           ├── name: :foo
-        │   │   │           ├── message_loc: (33,0)-(33,3) = "foo"
-        │   │   │           ├── opening_loc: ∅
-        │   │   │           ├── arguments: ∅
-        │   │   │           ├── closing_loc: ∅
-        │   │   │           └── block: ∅
-        │   │   ├── consequent: ∅
-        │   │   └── end_keyword_loc: ∅
-        │   ├── keyword_loc: (33,11)-(33,17) = "rescue"
-        │   └── rescue_expression:
-        │       @ CallNode (location: (33,18)-(33,21))
-        │       ├── flags: variable_call, ignore_visibility
-        │       ├── receiver: ∅
-        │       ├── call_operator_loc: ∅
-        │       ├── name: :baz
-        │       ├── message_loc: (33,18)-(33,21) = "baz"
-        │       ├── opening_loc: ∅
-        │       ├── arguments: ∅
-        │       ├── closing_loc: ∅
-        │       └── block: ∅
-        └── @ LocalVariableWriteNode (location: (35,0)-(35,18))
-            ├── name: :z
-            ├── depth: 0
-            ├── name_loc: (35,0)-(35,1) = "z"
-            ├── value:
-            │   @ RescueModifierNode (location: (35,4)-(35,18))
-            │   ├── expression:
-            │   │   @ CallNode (location: (35,4)-(35,7))
-            │   │   ├── flags: ignore_visibility
+        └── @ RescueModifierNode (location: (33,0)-(33,21))
+            ├── expression:
+            │   @ IfNode (location: (33,0)-(33,10))
+            │   ├── if_keyword_loc: (33,4)-(33,6) = "if"
+            │   ├── predicate:
+            │   │   @ CallNode (location: (33,7)-(33,10))
+            │   │   ├── flags: variable_call, ignore_visibility
             │   │   ├── receiver: ∅
             │   │   ├── call_operator_loc: ∅
-            │   │   ├── name: :x
-            │   │   ├── message_loc: (35,4)-(35,5) = "x"
+            │   │   ├── name: :bar
+            │   │   ├── message_loc: (33,7)-(33,10) = "bar"
             │   │   ├── opening_loc: ∅
-            │   │   ├── arguments:
-            │   │   │   @ ArgumentsNode (location: (35,6)-(35,7))
-            │   │   │   ├── flags: ∅
-            │   │   │   └── arguments: (length: 1)
-            │   │   │       └── @ CallNode (location: (35,6)-(35,7))
-            │   │   │           ├── flags: variable_call, ignore_visibility
-            │   │   │           ├── receiver: ∅
-            │   │   │           ├── call_operator_loc: ∅
-            │   │   │           ├── name: :y
-            │   │   │           ├── message_loc: (35,6)-(35,7) = "y"
-            │   │   │           ├── opening_loc: ∅
-            │   │   │           ├── arguments: ∅
-            │   │   │           ├── closing_loc: ∅
-            │   │   │           └── block: ∅
+            │   │   ├── arguments: ∅
             │   │   ├── closing_loc: ∅
             │   │   └── block: ∅
-            │   ├── keyword_loc: (35,8)-(35,14) = "rescue"
-            │   └── rescue_expression:
-            │       @ CallNode (location: (35,15)-(35,18))
-            │       ├── flags: ignore_visibility
-            │       ├── receiver: ∅
-            │       ├── call_operator_loc: ∅
-            │       ├── name: :c
-            │       ├── message_loc: (35,15)-(35,16) = "c"
-            │       ├── opening_loc: ∅
-            │       ├── arguments:
-            │       │   @ ArgumentsNode (location: (35,17)-(35,18))
-            │       │   ├── flags: ∅
-            │       │   └── arguments: (length: 1)
-            │       │       └── @ CallNode (location: (35,17)-(35,18))
-            │       │           ├── flags: variable_call, ignore_visibility
-            │       │           ├── receiver: ∅
-            │       │           ├── call_operator_loc: ∅
-            │       │           ├── name: :d
-            │       │           ├── message_loc: (35,17)-(35,18) = "d"
-            │       │           ├── opening_loc: ∅
-            │       │           ├── arguments: ∅
-            │       │           ├── closing_loc: ∅
-            │       │           └── block: ∅
-            │       ├── closing_loc: ∅
-            │       └── block: ∅
-            └── operator_loc: (35,2)-(35,3) = "="
+            │   ├── then_keyword_loc: ∅
+            │   ├── statements:
+            │   │   @ StatementsNode (location: (33,0)-(33,3))
+            │   │   └── body: (length: 1)
+            │   │       └── @ CallNode (location: (33,0)-(33,3))
+            │   │           ├── flags: variable_call, ignore_visibility
+            │   │           ├── receiver: ∅
+            │   │           ├── call_operator_loc: ∅
+            │   │           ├── name: :foo
+            │   │           ├── message_loc: (33,0)-(33,3) = "foo"
+            │   │           ├── opening_loc: ∅
+            │   │           ├── arguments: ∅
+            │   │           ├── closing_loc: ∅
+            │   │           └── block: ∅
+            │   ├── consequent: ∅
+            │   └── end_keyword_loc: ∅
+            ├── keyword_loc: (33,11)-(33,17) = "rescue"
+            └── rescue_expression:
+                @ CallNode (location: (33,18)-(33,21))
+                ├── flags: variable_call, ignore_visibility
+                ├── receiver: ∅
+                ├── call_operator_loc: ∅
+                ├── name: :baz
+                ├── message_loc: (33,18)-(33,21) = "baz"
+                ├── opening_loc: ∅
+                ├── arguments: ∅
+                ├── closing_loc: ∅
+                └── block: ∅