Merge pull request #2983 from ruby/shebang-callback

Callback on shebang switches

include/prism/options.h

@@ -40,6 +40,23 @@ typedef struct pm_options_scope {
     pm_string_t *locals;
 } pm_options_scope_t;
 
+// Forward declaration needed by the callback typedef.
+struct pm_options;
+
+/**
+ * The callback called when additional switches are found in a shebang comment
+ * that need to be processed by the runtime.
+ *
+ * @param options The options struct that may be updated by this callback.
+ *   Certain fields will be checked for changes, specifically encoding,
+ *   command_line, and frozen_string_literal.
+ * @param source The source of the shebang comment.
+ * @param length The length of the source.
+ * @param shebang_callback_data Any additional data that should be passed along
+ *   to the callback.
+ */
+typedef void (*pm_options_shebang_callback_t)(struct pm_options *options, const uint8_t *source, size_t length, void *shebang_callback_data);
+
 /**
  * The version of Ruby syntax that we should be parsing with. This is used to
  * allow consumers to specify which behavior they want in case they need to
@@ -56,7 +73,19 @@ typedef enum {
 /**
  * The options that can be passed to the parser.
  */
-typedef struct {
+typedef struct pm_options {
+    /**
+     * The callback to call when additional switches are found in a shebang
+     * comment.
+     */
+    pm_options_shebang_callback_t shebang_callback;
+
+    /**
+     * Any additional data that should be passed along to the shebang callback
+     * if one was set.
+     */
+    void *shebang_callback_data;
+
     /** The name of the file that is currently being parsed. */
     pm_string_t filepath;
 
@@ -149,6 +178,16 @@ static const uint8_t PM_OPTIONS_COMMAND_LINE_P = 0x10;
  */
 static const uint8_t PM_OPTIONS_COMMAND_LINE_X = 0x20;
 
+/**
+ * Set the shebang callback option on the given options struct.
+ *
+ * @param options The options struct to set the shebang callback on.
+ * @param shebang_callback The shebang callback to set.
+ * @param shebang_callback_data Any additional data that should be passed along
+ *   to the callback.
+ */
+PRISM_EXPORTED_FUNCTION void pm_options_shebang_callback_set(pm_options_t *options, pm_options_shebang_callback_t shebang_callback, void *shebang_callback_data);
+
 /**
  * Set the filepath option on the given options struct.
  *

src/options.c

@@ -1,5 +1,14 @@
 #include "prism/options.h"
 
+/**
+ * Set the shebang callback option on the given options struct.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_options_shebang_callback_set(pm_options_t *options, pm_options_shebang_callback_t shebang_callback, void *shebang_callback_data) {
+    options->shebang_callback = shebang_callback;
+    options->shebang_callback_data = shebang_callback_data;
+}
+
 /**
  * Set the filepath option on the given options struct.
  */

src/prism.c

@@ -21713,6 +21713,33 @@ pm_parser_warn_shebang_carriage_return(pm_parser_t *parser, const uint8_t *start
     }
 }
 
+/**
+ * Process the shebang when initializing the parser. This function assumes that
+ * the shebang_callback option has already been checked for nullability.
+ */
+static void
+pm_parser_init_shebang(pm_parser_t *parser, const pm_options_t *options, const char *engine, size_t length) {
+    const char *switches = pm_strnstr(engine, " -", length);
+    if (switches == NULL) return;
+
+    pm_options_t next_options = *options;
+    options->shebang_callback(
+        &next_options,
+        (const uint8_t *) (switches + 1),
+        length - ((size_t) (switches - engine)) - 1,
+        options->shebang_callback_data
+    );
+
+    size_t encoding_length;
+    if ((encoding_length = pm_string_length(&next_options.encoding)) > 0) {
+        const uint8_t *encoding_source = pm_string_source(&next_options.encoding);
+        parser_lex_magic_comment_encoding_value(parser, encoding_source, encoding_source + encoding_length);
+    }
+
+    parser->command_line = next_options.command_line;
+    parser->frozen_string_literal = next_options.frozen_string_literal;
+}
+
 /**
  * Initialize a parser with the given start and end pointers.
  */
@@ -21872,9 +21899,13 @@ pm_parser_init(pm_parser_t *parser, const uint8_t *source, size_t size, const pm
         const uint8_t *newline = next_newline(parser->start, parser->end - parser->start);
         size_t length = (size_t) ((newline != NULL ? newline : parser->end) - parser->start);
 
-        if (pm_strnstr((const char *) parser->start, "ruby", length) != NULL) {
+        const char *engine;
+        if ((engine = pm_strnstr((const char *) parser->start, "ruby", length)) != NULL) {
             pm_parser_warn_shebang_carriage_return(parser, parser->start, length);
             if (newline != NULL) parser->encoding_comment_start = newline + 1;
+            if (options != NULL && options->shebang_callback != NULL) {
+                pm_parser_init_shebang(parser, options, engine, length - ((size_t) (engine - (const char *) parser->start)));
+            }
             search_shebang = false;
         } else if (!parser->parsing_eval) {
             search_shebang = true;
@@ -21908,9 +21939,13 @@ pm_parser_init(pm_parser_t *parser, const uint8_t *source, size_t size, const pm
                     pm_parser_warn_shebang_carriage_return(parser, cursor, length);
                 }
 
-                if (pm_strnstr((const char *) cursor, "ruby", length) != NULL) {
+                const char *engine;
+                if ((engine = pm_strnstr((const char *) cursor, "ruby", length)) != NULL) {
                     found_shebang = true;
-                    parser->encoding_comment_start = newline + 1;
+                    if (newline != NULL) parser->encoding_comment_start = newline + 1;
+                    if (options != NULL && options->shebang_callback != NULL) {
+                        pm_parser_init_shebang(parser, options, engine, length - ((size_t) (engine - (const char *) cursor)));
+                    }
                     break;
                 }
             }