Make it possible to lazily deserialize DefNode in Loader.java

templates/java/org/ruby_lang/prism/Loader.java.erb

@@ -16,7 +16,7 @@ import java.util.Locale;
 public class Loader {
 
     public static ParseResult load(byte[] serialized, byte[] sourceBytes) {
-        return new Loader(serialized, sourceBytes).load();
+        return new Loader(serialized).load(sourceBytes);
     }
 
     // Overridable methods
@@ -71,19 +71,21 @@ public class Loader {
     }
 
     private final ByteBuffer buffer;
-    private final Nodes.Source source;
     protected String encodingName;
     <%- if string_type == "String" -%>
     private Charset encodingCharset;
     <%- end -%>
     private ConstantPool constantPool;
 
-    protected Loader(byte[] serialized, byte[] sourceBytes) {
+    protected Loader(byte[] serialized) {
         this.buffer = ByteBuffer.wrap(serialized).order(ByteOrder.nativeOrder());
-        this.source = new Nodes.Source(sourceBytes);
     }
 
-    protected ParseResult load() {
+    // We pass sourceBytes here and not in the constructor to avoid keeping
+    // the sourceBytes in memory unnecessarily with lazy DefNode's which hold on the Loader.
+    protected ParseResult load(byte[] sourceBytes) {
+        Nodes.Source source = new Nodes.Source(sourceBytes);
+
         expect((byte) 'P', "incorrect prism header");
         expect((byte) 'R', "incorrect prism header");
         expect((byte) 'I', "incorrect prism header");
@@ -331,6 +333,10 @@ public class Loader {
         return negative ? result.negate() : result;
     }
 
+    <%-
+    base_params = [*("nodeId" if Prism::Template::INCLUDE_NODE_ID), "startOffset", "length"]
+    base_params_sig = base_params.map { "int #{_1}" }.join(", ")
+    -%>
     private Nodes.Node loadNode() {
         int type = buffer.get() & 0xFF;
         <%- if Prism::Template::INCLUDE_NODE_ID -%>
@@ -347,7 +353,7 @@ public class Loader {
             params = []
             params << "nodeId" if Prism::Template::INCLUDE_NODE_ID
             params << "startOffset" << "length"
-            params << "buffer.getInt()" if node.needs_serialized_length?
+            params << "buffer.getInt()" << "null" if node.needs_serialized_length?
             params << "loadFlags()" if node.flags
             params.concat node.semantic_fields.map { |field|
               case field
@@ -370,13 +376,44 @@ public class Loader {
               else raise
               end
             }
+            $DefNode_params = params if node.name == "DefNode"
             -%>
+                <%- if node.name == "DefNode" -%>
+                return loadDefNode(<%= base_params.join(", ") -%>);
+                <%- else -%>
                 return new Nodes.<%= node.name %>(<%= params.join(", ") -%>);
+                <%- end -%>
             <%- end -%>
             default:
                 throw new Error("Unknown node type: " + type);
         }
     }
+
+    // Can be overridden to use createLazyDefNode instead
+    protected Nodes.DefNode loadDefNode(<%= base_params_sig -%>) {
+        return createDefNode(<%= base_params.join(", ") -%>);
+    }
+
+    protected Nodes.DefNode createLazyDefNode(<%= base_params_sig -%>) {
+        int bufferPosition = buffer.position();
+        int serializedLength = buffer.getInt();
+        // Load everything except the body and locals, because the name, receiver, parameters are still needed for lazily defining the method
+        Nodes.DefNode lazyDefNode = new Nodes.DefNode(<%= base_params.join(", ") -%>, -bufferPosition, this, loadConstant(), loadOptionalNode(), (Nodes.ParametersNode) loadOptionalNode(), null, Nodes.EMPTY_STRING_ARRAY);
+        buffer.position(bufferPosition + serializedLength); // skip past the serialized DefNode
+        return lazyDefNode;
+    }
+
+    protected Nodes.DefNode createDefNode(<%= base_params_sig -%>) {
+        return new Nodes.DefNode(<%= $DefNode_params.join(", ") -%>);
+    }
+
+    Nodes.DefNode createDefNodeFromSavedPosition(<%= base_params_sig -%>, int bufferPosition) {
+        // This method mutates the buffer position and may be called from different threads so we must synchronize
+        synchronized (this) {
+            buffer.position(bufferPosition);
+            return createDefNode(<%= base_params.join(", ") -%>);
+        }
+    }
     <%- array_types.uniq.each do |type| -%>
 
     private static final Nodes.<%= type %>[] EMPTY_<%= type %>_ARRAY = {};

templates/java/org/ruby_lang/prism/Nodes.java.erb

@@ -142,7 +142,7 @@ public abstract class Nodes {
 
         protected abstract String toString(String indent);
     }
-<%# FLAGS -%>
+<%-# FLAGS -%>
     <%- flags.each do |flag| -%>
 
     /**
@@ -194,7 +194,7 @@ public abstract class Nodes {
         <%- end -%>
     }
 <%- end -%>
-<%# NODES -%>
+<%-# NODES -%>
     <%- nodes.each do |node| -%>
 
     /**
@@ -207,6 +207,7 @@ public abstract class Nodes {
     public static final class <%= node.name -%> extends Node {
         <%- if node.needs_serialized_length? -%>
         public final int serializedLength;
+        public final Loader loader;
         <%- end -%>
         <%- if node.flags -%>
         public final short flags;
@@ -234,7 +235,10 @@ public abstract class Nodes {
           params = []
           params << "int nodeId" if Prism::Template::INCLUDE_NODE_ID
           params << "int startOffset" << "int length"
-          params << "int serializedLength" if node.needs_serialized_length?
+          if node.needs_serialized_length?
+            params << "int serializedLength"
+            params << "Loader loader"
+        end
           params << "short flags" if node.flags
           params.concat(node.semantic_fields.map { |field| "#{field.java_type} #{field.name}" })
         -%>
@@ -246,6 +250,7 @@ public abstract class Nodes {
         <%- end -%>
         <%- if node.needs_serialized_length? -%>
             this.serializedLength = serializedLength;
+            this.loader = loader;
         <%- end -%>
         <%- if node.flags -%>
             this.flags = flags;
@@ -254,7 +259,22 @@ public abstract class Nodes {
             this.<%= field.name %> = <%= field.name %>;
         <%- end -%>
         }
-        <%# methods for flags -%>
+        <%-# extra methods for DefNode -%>
+        <%- if node.needs_serialized_length? -%>
+
+        public boolean isLazy() {
+            return serializedLength < 0;
+        }
+
+        public <%= node.name -%> getNonLazy() {
+            if (isLazy()) {
+                return loader.createDefNodeFromSavedPosition(<%= "nodeId, " if Prism::Template::INCLUDE_NODE_ID %>startOffset, length, -serializedLength);
+            } else {
+                return this;
+            }
+        }
+        <%- end -%>
+        <%-# methods for flags -%>
         <%- if (node_flags = node.flags) -%>
         <%- node_flags.values.each do |value| -%>
 
@@ -263,7 +283,7 @@ public abstract class Nodes {
         }
         <%- end -%>
         <%- end -%>
-        <%# potential override of setNewLineFlag() -%>
+        <%-# potential override of setNewLineFlag() -%>
         <%- if node.newline == false -%>
 
         @Override

templates/src/serialize.c.erb

@@ -50,8 +50,6 @@ static void
 pm_serialize_node(pm_parser_t *parser, pm_node_t *node, pm_buffer_t *buffer) {
     pm_buffer_append_byte(buffer, (uint8_t) PM_NODE_TYPE(node));
 
-    size_t offset = buffer->length;
-
     <%- if Prism::Template::INCLUDE_NODE_ID -%>
     pm_buffer_append_varuint(buffer, node->node_id);
     <%- end -%>
@@ -126,7 +124,7 @@ pm_serialize_node(pm_parser_t *parser, pm_node_t *node, pm_buffer_t *buffer) {
             <%- end -%>
             <%- if node.needs_serialized_length? -%>
             // serialize length
-            uint32_t length = pm_sizet_to_u32(buffer->length - offset - sizeof(uint32_t));
+            uint32_t length = pm_sizet_to_u32(buffer->length - length_offset);
             memcpy(buffer->value + length_offset, &length, sizeof(uint32_t));
             <%- end -%>
             break;