[ruby/prism] Support -p, -n, -a, and -l command line options

https://github.com/ruby/prism/commit/959eb506ca

lib/prism/ffi.rb

@@ -342,6 +342,18 @@ module Prism
       template << "C"
       values << (options.fetch(:frozen_string_literal, false) ? 1 : 0)
 
+      template << "C"
+      values << (options.fetch(:command_line_p, false) ? 1 : 0)
+
+      template << "C"
+      values << (options.fetch(:command_line_n, false) ? 1 : 0)
+
+      template << "C"
+      values << (options.fetch(:command_line_l, false) ? 1 : 0)
+
+      template << "C"
+      values << (options.fetch(:command_line_a, false) ? 1 : 0)
+
       template << "C"
       values << { nil => 0, "3.3.0" => 1, "3.4.0" => 0, "latest" => 0 }.fetch(options[:version])
 

prism/extension.c

@@ -29,6 +29,10 @@ ID rb_option_id_line;
 ID rb_option_id_frozen_string_literal;
 ID rb_option_id_version;
 ID rb_option_id_scopes;
+ID rb_option_id_command_line_p;
+ID rb_option_id_command_line_n;
+ID rb_option_id_command_line_l;
+ID rb_option_id_command_line_a;
 
 /******************************************************************************/
 /* IO of Ruby code                                                            */
@@ -149,6 +153,14 @@ build_options_i(VALUE key, VALUE value, VALUE argument) {
         }
     } else if (key_id == rb_option_id_scopes) {
         if (!NIL_P(value)) build_options_scopes(options, value);
+    } else if (key_id == rb_option_id_command_line_p) {
+        if (!NIL_P(value)) pm_options_command_line_p_set(options, value == Qtrue);
+    } else if (key_id == rb_option_id_command_line_n) {
+        if (!NIL_P(value)) pm_options_command_line_n_set(options, value == Qtrue);
+    } else if (key_id == rb_option_id_command_line_l) {
+        if (!NIL_P(value)) pm_options_command_line_l_set(options, value == Qtrue);
+    } else if (key_id == rb_option_id_command_line_a) {
+        if (!NIL_P(value)) pm_options_command_line_a_set(options, value == Qtrue);
     } else {
         rb_raise(rb_eArgError, "unknown keyword: %"PRIsVALUE, key);
     }
@@ -1232,6 +1244,10 @@ Init_prism(void) {
     rb_option_id_frozen_string_literal = rb_intern_const("frozen_string_literal");
     rb_option_id_version = rb_intern_const("version");
     rb_option_id_scopes = rb_intern_const("scopes");
+    rb_option_id_command_line_p = rb_intern_const("command_line_p");
+    rb_option_id_command_line_n = rb_intern_const("command_line_n");
+    rb_option_id_command_line_l = rb_intern_const("command_line_l");
+    rb_option_id_command_line_a = rb_intern_const("command_line_a");
 
     /**
      * The version of the prism library.

prism/node.h

@@ -29,6 +29,15 @@ bool pm_node_list_grow(pm_node_list_t *list);
  */
 void pm_node_list_append(pm_node_list_t *list, pm_node_t *node);
 
+/**
+ * Prepend a new node onto the beginning of the node list.
+ *
+ * @param list The list to prepend to.
+ * @param node The node to prepend.
+ */
+void
+pm_node_list_prepend(pm_node_list_t *list, pm_node_t *node);
+
 /**
  * Free the internal memory associated with the given node list.
  *

prism/options.c

@@ -32,6 +32,38 @@ pm_options_frozen_string_literal_set(pm_options_t *options, bool frozen_string_l
     options->frozen_string_literal = frozen_string_literal;
 }
 
+/**
+ * Sets the -p command line option on the given options struct.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_options_command_line_p_set(pm_options_t *options, bool command_line_p) {
+    options->command_line_p = command_line_p;
+}
+
+/**
+ * Sets the -n command line option on the given options struct.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_options_command_line_n_set(pm_options_t *options, bool command_line_n) {
+    options->command_line_n = command_line_n;
+}
+
+/**
+ * Sets the -l command line option on the given options struct.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_options_command_line_l_set(pm_options_t *options, bool command_line_l) {
+    options->command_line_l = command_line_l;
+}
+
+/**
+ * Sets the -a command line option on the given options struct.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_options_command_line_a_set(pm_options_t *options, bool command_line_a) {
+    options->command_line_a = command_line_a;
+}
+
 /**
  * Set the version option on the given options struct by parsing the given
  * string. If the string contains an invalid option, this returns false.
@@ -193,7 +225,11 @@ pm_options_read(pm_options_t *options, const char *data) {
         data += encoding_length;
     }
 
-    options->frozen_string_literal = *data++;
+    options->frozen_string_literal = (*data++) ? true : false;
+    options->command_line_p = (*data++) ? true : false;
+    options->command_line_n = (*data++) ? true : false;
+    options->command_line_l = (*data++) ? true : false;
+    options->command_line_a = (*data++) ? true : false;
     options->version = (pm_options_version_t) *data++;
 
     uint32_t scopes_count = pm_options_read_u32(data);

prism/options.h

@@ -78,6 +78,18 @@ typedef struct {
 
     /** Whether or not the frozen string literal option has been set. */
     bool frozen_string_literal;
+
+    /** Whether or not the -p command line option has been set. */
+    bool command_line_p;
+
+    /** Whether or not the -n command line option has been set. */
+    bool command_line_n;
+
+    /** Whether or not the -l command line option has been set. */
+    bool command_line_l;
+
+    /** Whether or not the -a command line option has been set. */
+    bool command_line_a;
 } pm_options_t;
 
 /**
@@ -112,6 +124,38 @@ PRISM_EXPORTED_FUNCTION void pm_options_encoding_set(pm_options_t *options, cons
  */
 PRISM_EXPORTED_FUNCTION void pm_options_frozen_string_literal_set(pm_options_t *options, bool frozen_string_literal);
 
+/**
+ * Sets the -p command line option on the given options struct.
+ *
+ * @param options The options struct to set the -p command line option on.
+ * @param command_line_p The -p command line option to set.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_command_line_p_set(pm_options_t *options, bool command_line_p);
+
+/**
+ * Sets the -n command line option on the given options struct.
+ *
+ * @param options The options struct to set the -n command line option on.
+ * @param command_line_n The -n command line option to set.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_command_line_n_set(pm_options_t *options, bool command_line_n);
+
+/**
+ * Sets the -l command line option on the given options struct.
+ *
+ * @param options The options struct to set the -l command line option on.
+ * @param command_line_l The -l command line option to set.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_command_line_l_set(pm_options_t *options, bool command_line_l);
+
+/**
+ * Sets the -a command line option on the given options struct.
+ *
+ * @param options The options struct to set the -a command line option on.
+ * @param command_line_a The -a command line option to set.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_command_line_a_set(pm_options_t *options, bool command_line_a);
+
 /**
  * Set the version option on the given options struct by parsing the given
  * string. If the string contains an invalid option, this returns false.
@@ -186,7 +230,10 @@ PRISM_EXPORTED_FUNCTION void pm_options_free(pm_options_t *options);
  * | `4`     | the length the encoding    |
  * | ...     | the encoding bytes         |
  * | `1`     | frozen string literal      |
- * | `1`     | suppress warnings          |
+ * | `1`     | -p command line option     |
+ * | `1`     | -n command line option     |
+ * | `1`     | -l command line option     |
+ * | `1`     | -a command line option     |
  * | `1`     | the version                |
  * | `4`     | the number of scopes       |
  * | ...     | the scopes                 |

prism/parser.h

@@ -736,6 +736,30 @@ struct pm_parser {
      * a true value.
      */
     bool frozen_string_literal;
+
+    /**
+     * Whether or not -p was present on the command line that invoked the
+     * parser. -p prints the value of $_ at the end of each loop.
+     */
+    bool command_line_p;
+
+    /**
+     * Whether or not -n was present on the command line that invoked the
+     * parser. -n wraps the script in a while gets loop.
+     */
+    bool command_line_n;
+
+    /**
+     * Whether or not -l was present on the command line that invoked the
+     * parser. -l chomps the input line by default.
+     */
+    bool command_line_l;
+
+    /**
+     * Whether or not -a was present on the command line that invoked the
+     * parser. -a splits the input line $_ into $F.
+     */
+    bool command_line_a;
 };
 
 #endif

prism/prism.c

@@ -1904,6 +1904,24 @@ pm_call_node_call_create(pm_parser_t *parser, pm_node_t *receiver, pm_token_t *o
     return node;
 }
 
+/**
+ * Allocate and initialize a new synthesized CallNode node from a call expression.
+ */
+static pm_call_node_t *
+pm_call_node_call_synthesized_create(pm_parser_t *parser, pm_node_t *receiver, const char *message, pm_arguments_node_t *arguments) {
+    pm_call_node_t *node = pm_call_node_create(parser, 0);
+    node->base.location.start = parser->start;
+    node->base.location.end = parser->end;
+
+    node->receiver = receiver;
+    node->call_operator_loc = (pm_location_t) { .start = NULL, .end = NULL };
+    node->message_loc = (pm_location_t) { .start = NULL, .end = NULL };
+    node->arguments = arguments;
+
+    node->name = pm_parser_constant_id_constant(parser, message, strlen(message));
+    return node;
+}
+
 /**
  * Allocate and initialize a new CallNode node from a call to a method name
  * without a receiver that could not have been a local variable read.
@@ -1925,6 +1943,22 @@ pm_call_node_fcall_create(pm_parser_t *parser, pm_token_t *message, pm_arguments
     return node;
 }
 
+/**
+ * Allocate and initialize a new CallNode node from a synthesized call to a
+ * method name with the given arguments.
+ */
+static pm_call_node_t *
+pm_call_node_fcall_synthesized_create(pm_parser_t *parser, pm_arguments_node_t *arguments, pm_constant_id_t name) {
+    pm_call_node_t *node = pm_call_node_create(parser, PM_CALL_NODE_FLAGS_IGNORE_VISIBILITY);
+
+    node->base.location.start = parser->start;
+    node->base.location.end = parser->start;
+    node->arguments = arguments;
+
+    node->name = name;
+    return node;
+}
+
 /**
  * Allocate and initialize a new CallNode node from a not expression.
  */
@@ -3568,7 +3602,25 @@ pm_global_variable_read_node_create(pm_parser_t *parser, const pm_token_t *name)
 }
 
 /**
- * Allocate a new GlobalVariableWriteNode node.
+ * Allocate and initialize a new synthesized GlobalVariableReadNode node.
+ */
+static pm_global_variable_read_node_t *
+pm_global_variable_read_node_synthesized_create(pm_parser_t *parser, pm_constant_id_t name) {
+    pm_global_variable_read_node_t *node = PM_ALLOC_NODE(parser, pm_global_variable_read_node_t);
+
+    *node = (pm_global_variable_read_node_t) {
+        {
+            .type = PM_GLOBAL_VARIABLE_READ_NODE,
+            .location = { .start = parser->start, .end = parser->start }
+        },
+        .name = name
+    };
+
+    return node;
+}
+
+/**
+ * Allocate and initialize a new GlobalVariableWriteNode node.
  */
 static pm_global_variable_write_node_t *
 pm_global_variable_write_node_create(pm_parser_t *parser, pm_node_t *target, const pm_token_t *operator, pm_node_t *value) {
@@ -3591,6 +3643,27 @@ pm_global_variable_write_node_create(pm_parser_t *parser, pm_node_t *target, con
     return node;
 }
 
+/**
+ * Allocate and initialize a new synthesized GlobalVariableWriteNode node.
+ */
+static pm_global_variable_write_node_t *
+pm_global_variable_write_node_synthesized_create(pm_parser_t *parser, pm_constant_id_t name, pm_node_t *value) {
+    pm_global_variable_write_node_t *node = PM_ALLOC_NODE(parser, pm_global_variable_write_node_t);
+
+    *node = (pm_global_variable_write_node_t) {
+        {
+            .type = PM_GLOBAL_VARIABLE_WRITE_NODE,
+            .location = { .start = parser->start, .end = parser->start }
+        },
+        .name = name,
+        .name_loc = { .start = parser->start, .end = parser->start },
+        .operator_loc = { .start = parser->start, .end = parser->start },
+        .value = value
+    };
+
+    return node;
+}
+
 /**
  * Allocate a new HashNode node.
  */
@@ -5677,21 +5750,38 @@ pm_statements_node_location_set(pm_statements_node_t *node, const uint8_t *start
     node->base.location = (pm_location_t) { .start = start, .end = end };
 }
 
+/**
+ * Update the location of the statements node based on the statement that is
+ * being added to the list.
+ */
+static inline void
+pm_statements_node_body_update(pm_statements_node_t *node, pm_node_t *statement) {
+    if (pm_statements_node_body_length(node) == 0 || statement->location.start < node->base.location.start) {
+        node->base.location.start = statement->location.start;
+    }
+
+    if (statement->location.end > node->base.location.end) {
+        node->base.location.end = statement->location.end;
+    }
+}
+
 /**
  * Append a new node to the given StatementsNode node's body.
  */
 static void
 pm_statements_node_body_append(pm_statements_node_t *node, pm_node_t *statement) {
-    if (pm_statements_node_body_length(node) == 0 || statement->location.start < node->base.location.start) {
+    pm_statements_node_body_update(node, statement);
-        node->base.location.start = statement->location.start;
-    }
-    if (statement->location.end > node->base.location.end) {
-        node->base.location.end = statement->location.end;
-    }
-
     pm_node_list_append(&node->body, statement);
+    pm_node_flag_set(statement, PM_NODE_FLAG_NEWLINE);
+}
 
-    // Every statement gets marked as a place where a newline can occur.
+/**
+ * Prepend a new node to the given StatementsNode node's body.
+ */
+static void
+pm_statements_node_body_prepend(pm_statements_node_t *node, pm_node_t *statement) {
+    pm_statements_node_body_update(node, statement);
+    pm_node_list_prepend(&node->body, statement);
     pm_node_flag_set(statement, PM_NODE_FLAG_NEWLINE);
 }
 
@@ -5898,6 +5988,27 @@ pm_symbol_node_label_create(pm_parser_t *parser, const pm_token_t *token) {
     return node;
 }
 
+/**
+ * Allocate and initialize a new synthesized SymbolNode node.
+ */
+static pm_symbol_node_t *
+pm_symbol_node_synthesized_create(pm_parser_t *parser, const char *content) {
+    pm_symbol_node_t *node = PM_ALLOC_NODE(parser, pm_symbol_node_t);
+
+    *node = (pm_symbol_node_t) {
+        {
+            .type = PM_SYMBOL_NODE,
+            .flags = PM_NODE_FLAG_STATIC_LITERAL | PM_SYMBOL_FLAGS_FORCED_US_ASCII_ENCODING,
+            .location = { .start = parser->start, .end = parser->start }
+        },
+        .value_loc = { .start = parser->start, .end = parser->start },
+        .unescaped = { 0 }
+    };
+
+    pm_string_constant_init(&node->unescaped, content, strlen(content));
+    return node;
+}
+
 /**
  * Check if the given node is a label in a hash.
  */
@@ -6000,6 +6111,22 @@ pm_true_node_create(pm_parser_t *parser, const pm_token_t *token) {
     return node;
 }
 
+/**
+ * Allocate and initialize a new synthesized TrueNode node.
+ */
+static pm_true_node_t *
+pm_true_node_synthesized_create(pm_parser_t *parser) {
+    pm_true_node_t *node = PM_ALLOC_NODE(parser, pm_true_node_t);
+
+    *node = (pm_true_node_t) {{
+        .type = PM_TRUE_NODE,
+        .flags = PM_NODE_FLAG_STATIC_LITERAL,
+        .location = { .start = parser->start, .end = parser->end }
+    }};
+
+    return node;
+}
+
 /**
  * Allocate and initialize a new UndefNode node.
  */
@@ -6251,6 +6378,27 @@ pm_while_node_modifier_create(pm_parser_t *parser, const pm_token_t *keyword, pm
     return node;
 }
 
+/**
+ * Allocate and initialize a new synthesized while loop.
+ */
+static pm_while_node_t *
+pm_while_node_synthesized_create(pm_parser_t *parser, pm_node_t *predicate, pm_statements_node_t *statements) {
+    pm_while_node_t *node = PM_ALLOC_NODE(parser, pm_while_node_t);
+
+    *node = (pm_while_node_t) {
+        {
+            .type = PM_WHILE_NODE,
+            .location = { .start = parser->start, .end = parser->start }
+        },
+        .keyword_loc = { .start = parser->start, .end = parser->start },
+        .closing_loc = { .start = parser->start, .end = parser->start },
+        .predicate = predicate,
+        .statements = statements
+    };
+
+    return node;
+}
+
 /**
  * Allocate and initialize a new XStringNode node with the given unescaped
  * string.
@@ -18108,6 +18256,80 @@ parse_expression(pm_parser_t *parser, pm_binding_power_t binding_power, bool acc
     return node;
 }
 
+/**
+ * ruby -p, ruby -n, ruby -a, and ruby -l options will mutate the AST. We
+ * perform that mutation here.
+ */
+static pm_statements_node_t *
+wrap_statements(pm_parser_t *parser, pm_statements_node_t *statements) {
+    if (parser->command_line_p) {
+        pm_arguments_node_t *arguments = pm_arguments_node_create(parser);
+        pm_arguments_node_arguments_append(
+            arguments,
+            (pm_node_t *) pm_global_variable_read_node_synthesized_create(parser, pm_parser_constant_id_constant(parser, "$_", 2))
+        );
+
+        pm_statements_node_body_append(statements, (pm_node_t *) pm_call_node_fcall_synthesized_create(
+            parser,
+            arguments,
+            pm_parser_constant_id_constant(parser, "print", 5)
+        ));
+    }
+
+    if (parser->command_line_n) {
+        if (parser->command_line_a) {
+            pm_arguments_node_t *arguments = pm_arguments_node_create(parser);
+            pm_arguments_node_arguments_append(
+                arguments,
+                (pm_node_t *) pm_global_variable_read_node_synthesized_create(parser, pm_parser_constant_id_constant(parser, "$;", 2))
+            );
+
+            pm_global_variable_read_node_t *receiver = pm_global_variable_read_node_synthesized_create(parser, pm_parser_constant_id_constant(parser, "$_", 2));
+            pm_call_node_t *call = pm_call_node_call_synthesized_create(parser, (pm_node_t *) receiver, "split", arguments);
+
+            pm_global_variable_write_node_t *write = pm_global_variable_write_node_synthesized_create(
+                parser,
+                pm_parser_constant_id_constant(parser, "$F", 2),
+                (pm_node_t *) call
+            );
+
+            pm_statements_node_body_prepend(statements, (pm_node_t *) write);
+        }
+
+        pm_arguments_node_t *arguments = pm_arguments_node_create(parser);
+        pm_arguments_node_arguments_append(
+            arguments,
+            (pm_node_t *) pm_global_variable_read_node_synthesized_create(parser, pm_parser_constant_id_constant(parser, "$/", 2))
+        );
+
+        if (parser->command_line_l) {
+            pm_keyword_hash_node_t *keywords = pm_keyword_hash_node_create(parser);
+            pm_keyword_hash_node_elements_append(keywords, (pm_node_t *) pm_assoc_node_create(
+                parser,
+                (pm_node_t *) pm_symbol_node_synthesized_create(parser, "chomp"),
+                &(pm_token_t) { .type = PM_TOKEN_NOT_PROVIDED, .start = parser->start, .end = parser->start },
+                (pm_node_t *) pm_true_node_synthesized_create(parser)
+            ));
+
+            pm_arguments_node_arguments_append(arguments, (pm_node_t *) keywords);
+        }
+
+        pm_statements_node_t *wrapped_statements = pm_statements_node_create(parser);
+        pm_statements_node_body_append(wrapped_statements, (pm_node_t *) pm_while_node_synthesized_create(
+            parser,
+            (pm_node_t *) pm_call_node_fcall_synthesized_create(parser, arguments, pm_parser_constant_id_constant(parser, "gets", 4)),
+            statements
+        ));
+
+        statements = wrapped_statements;
+    }
+
+    return statements;
+}
+
+/**
+ * Parse the top-level program node.
+ */
 static pm_node_t *
 parse_program(pm_parser_t *parser) {
     // If the current scope is NULL, then we want to push a new top level scope.
@@ -18132,6 +18354,12 @@ parse_program(pm_parser_t *parser) {
         pm_statements_node_location_set(statements, parser->start, parser->start);
     }
 
+    // At the top level, see if we need to wrap the statements in a program
+    // node with a while loop based on the options.
+    if (parser->command_line_p || parser->command_line_n) {
+        statements = wrap_statements(parser, statements);
+    }
+
     return (pm_node_t *) pm_program_node_create(parser, &locals, statements);
 }
 
@@ -18189,7 +18417,11 @@ pm_parser_init(pm_parser_t *parser, const uint8_t *source, size_t size, const pm
         .in_keyword_arg = false,
         .current_param_name = 0,
         .semantic_token_seen = false,
-        .frozen_string_literal = false
+        .frozen_string_literal = false,
+        .command_line_p = options != NULL && options->command_line_p,
+        .command_line_n = options != NULL && options->command_line_n,
+        .command_line_l = options != NULL && options->command_line_l,
+        .command_line_a = options != NULL && options->command_line_a
     };
 
     // Initialize the constant pool. We're going to completely guess as to the

prism/templates/src/node.c.erb

@@ -42,6 +42,18 @@ pm_node_list_append(pm_node_list_t *list, pm_node_t *node) {
     }
 }
 
+/**
+ * Prepend a new node onto the beginning of the node list.
+ */
+void
+pm_node_list_prepend(pm_node_list_t *list, pm_node_t *node) {
+    if (pm_node_list_grow(list)) {
+        memmove(list->nodes + 1, list->nodes, list->size * sizeof(pm_node_t *));
+        list->nodes[0] = node;
+        list->size++;
+    }
+}
+
 /**
  * Free the internal memory associated with the given node list.
  */

test/prism/command_line_test.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require_relative "test_helper"
+
+module Prism
+  class CommandLineTest < TestCase
+    def test_command_line_p
+      program = Prism.parse("1", command_line_p: true).value
+      statements = program.statements.body
+
+      assert_equal 2, statements.length
+      assert_kind_of CallNode, statements.last
+      assert_equal :print, statements.last.name
+    end
+
+    def test_command_line_n
+      program = Prism.parse("1", command_line_n: true).value
+      statements = program.statements.body
+
+      assert_equal 1, statements.length
+      assert_kind_of WhileNode, statements.first
+
+      predicate = statements.first.predicate
+      assert_kind_of CallNode, predicate
+      assert_equal :gets, predicate.name
+
+      arguments = predicate.arguments.arguments
+      assert_equal 1, arguments.length
+      assert_equal :$/, arguments.first.name
+    end
+
+    def test_command_line_a
+      program = Prism.parse("1", command_line_n: true, command_line_a: true).value
+      statements = program.statements.body
+
+      assert_equal 1, statements.length
+      assert_kind_of WhileNode, statements.first
+
+      statement = statements.first.statements.body.first
+      assert_kind_of GlobalVariableWriteNode, statement
+      assert_equal :$F, statement.name
+    end
+
+    def test_command_line_l
+      program = Prism.parse("1", command_line_n: true, command_line_l: true).value
+      statements = program.statements.body
+
+      assert_equal 1, statements.length
+      assert_kind_of WhileNode, statements.first
+
+      predicate = statements.first.predicate
+      assert_kind_of CallNode, predicate
+      assert_equal :gets, predicate.name
+
+      arguments = predicate.arguments.arguments
+      assert_equal 2, arguments.length
+      assert_equal :$/, arguments.first.name
+      assert_equal "chomp", arguments.last.elements.first.key.unescaped
+    end
+  end
+end