Add index and at APIs for constants

include/prism/constant_pool.h

@@ -11,6 +11,7 @@
 #define PRISM_CONSTANT_POOL_H
 
 #include "prism/compiler/exported.h"
+#include "prism/compiler/inline.h"
 #include "prism/compiler/nodiscard.h"
 #include "prism/compiler/nonnull.h"
 
@@ -28,6 +29,22 @@ extern "C" {
  */
 typedef uint32_t pm_constant_id_t;
 
+/**
+ * Returns the index of the constant with the given id among the constants of
+ * the parser that produced it. The indices run from 0 to one less than
+ * pm_parser_constants_size, so they can index a cache with one entry per
+ * constant that was filled using pm_parser_constant_at.
+ *
+ * @param constant_id the id of the constant. Optional constant fields on nodes
+ *     hold 0 when they are absent, which is not a valid id, so check for it
+ *     before calling this function.
+ * @returns the index of the constant
+ */
+static PRISM_INLINE size_t
+pm_constant_id_index(pm_constant_id_t constant_id) {
+    return (size_t) constant_id - 1;
+}
+
 /**
  * A list of constant IDs. Usually used to represent a set of locals.
  */

include/prism/parser.h

@@ -364,30 +364,38 @@ PRISM_EXPORTED_FUNCTION const pm_diagnostic_t * pm_parser_warnings_first(const p
 PRISM_EXPORTED_FUNCTION size_t pm_parser_constants_size(const pm_parser_t *parser) PRISM_NONNULL(1);
 
 /**
- * A callback function that can be used to process constants found while
- * parsing.
- */
-typedef void (*pm_constant_callback_t)(const pm_constant_t *constant, void *data);
-
-/**
- * Iterates over the constants in the constant pool associated with the given
- * parser and calls the given callback for each constant.
+ * Returns the constant at the given index in the constant pool associated with
+ * the given parser. Together with pm_constant_id_index, this lets a caller
+ * build a cache with one entry per constant and then find the entry for any
+ * constant id, as in:
+ *
+ * ```c
+ * size_t size = pm_parser_constants_size(parser);
+ * for (size_t index = 0; index < size; index++) {
+ *     cache[index] = convert(pm_parser_constant_at(parser, index));
+ * }
+ *
+ * value = cache[pm_constant_id_index(constant_id)];
+ * ```
  *
- * @param parser the parser whose constants we want to iterate over
- * @param callback the callback function to call for each constant. This function
- *     will be passed a pointer to the constant and the data parameter passed to
- *     this function.
- * @param data the data to pass to the callback function for each constant. This
- *     can be NULL if no data needs to be passed to the callback function.
+ * For any constant id, pm_parser_constant_at(parser,
+ * pm_constant_id_index(constant_id)) is pm_parser_constant(parser,
+ * constant_id).
+ *
+ * @param parser the parser whose constant pool we want to look up from
+ * @param index the index of the constant, which must be less than
+ *     pm_parser_constants_size(parser)
+ * @returns a pointer to the constant at the given index
  */
-PRISM_EXPORTED_FUNCTION void pm_parser_constants_each(const pm_parser_t *parser, pm_constant_callback_t callback, void *data) PRISM_NONNULL(1);
+PRISM_EXPORTED_FUNCTION const pm_constant_t * pm_parser_constant_at(const pm_parser_t *parser, size_t index) PRISM_NONNULL(1);
 
 /**
  * Returns a pointer to the constant at the given id in the constant pool
  * associated with the given parser.
  *
  * @param parser the parser whose constant pool we want to look up from
- * @param constant_id the id of the constant to look up (1-based)
+ * @param constant_id the id of the constant to look up. Optional constant
+ *     fields on nodes hold 0 when they are absent, which is not a valid id.
  * @returns a pointer to the constant at the given id
  */
 PRISM_EXPORTED_FUNCTION const pm_constant_t * pm_parser_constant(const pm_parser_t *parser, pm_constant_id_t constant_id) PRISM_NONNULL(1);

rust/ruby-prism-sys/build/main.rs

@@ -219,7 +219,7 @@ fn generate_bindings(ruby_include_path: &Path, clang_args: &[String], layout_tes
         .allowlist_function("pm_parser_comments_first")
         .allowlist_function("pm_parser_comments_size")
         .allowlist_function("pm_parser_constant")
-        .allowlist_function("pm_parser_constants_each")
+        .allowlist_function("pm_parser_constant_at")
         .allowlist_function("pm_parser_constants_size")
         .allowlist_function("pm_parser_data_loc")
         .allowlist_function("pm_parser_errors_first")

src/parser.c

@@ -7,6 +7,7 @@
 #include "prism/internal/magic_comments.h"
 #include "prism/internal/xxh3.h"
 
+#include <assert.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -338,15 +339,13 @@ pm_parser_constants_size(const pm_parser_t *parser) {
 }
 
 /**
- * Iterates over the constants in the constant pool associated with the given
- * parser and calls the given callback for each constant.
+ * Returns the constant at the given index in the constant pool associated with
+ * the given parser.
  */
-void
-pm_parser_constants_each(const pm_parser_t *parser, pm_constant_callback_t callback, void *data) {
-    for (uint32_t index = 0; index < parser->constant_pool.size; index++) {
-        const pm_constant_t *constant = &parser->constant_pool.constants[index];
-        callback(constant, data);
-    }
+const pm_constant_t *
+pm_parser_constant_at(const pm_parser_t *parser, size_t index) {
+    assert(index < parser->constant_pool.size);
+    return &parser->constant_pool.constants[index];
 }
 
 /**

templates/ext/prism/api_node.c.erb

@@ -127,32 +127,32 @@ pm_node_stack_pop(pm_node_stack_node_t **stack) {
     return visit;
 }
 
-typedef struct {
-    VALUE constants;
-    rb_encoding *encoding;
-} pm_ast_constants_each_data_t;
-
-static void
-pm_ast_constants_each(const pm_constant_t *constant, void *data) {
-    pm_ast_constants_each_data_t *constants_data = (pm_ast_constants_each_data_t *) data;
+/**
+ * Build the Symbol for the given constant.
+ */
+static VALUE
+pm_ast_constant(const pm_constant_t *constant, rb_encoding *encoding) {
     int state = 0;
 
-    VALUE string = rb_enc_str_new((const char *) pm_constant_start(constant), pm_constant_length(constant), constants_data->encoding);
+    VALUE string = rb_enc_str_new((const char *) pm_constant_start(constant), pm_constant_length(constant), encoding);
     VALUE value = rb_protect(rb_str_intern, string, &state);
 
     if (state != 0) {
         value = ID2SYM(rb_intern_const("?"));
         rb_set_errinfo(Qnil);
     }
 
-    rb_ary_push(constants_data->constants, value);
+    return value;
 }
 
 VALUE
 pm_ast_new(const pm_parser_t *parser, const pm_node_t *node, rb_encoding *encoding, VALUE source, bool freeze) {
-    VALUE constants = rb_ary_new_capa(pm_parser_constants_size(parser));
-    pm_ast_constants_each_data_t constants_data = { .constants = constants, .encoding = encoding };
-    pm_parser_constants_each(parser, pm_ast_constants_each, &constants_data);
+    size_t constants_size = pm_parser_constants_size(parser);
+    VALUE constants = rb_ary_new_capa((long) constants_size);
+
+    for (size_t index = 0; index < constants_size; index++) {
+        rb_ary_push(constants, pm_ast_constant(pm_parser_constant_at(parser, index), encoding));
+    }
 
     pm_arena_t *node_arena = pm_arena_new();
     pm_node_stack_node_t *node_stack = NULL;
@@ -237,15 +237,15 @@ pm_ast_new(const pm_parser_t *parser, const pm_node_t *node, rb_encoding *encodi
                     <%- when Prism::Template::ConstantField -%>
 #line <%= __LINE__ + 1 %> "prism/templates/ext/prism/<%= File.basename(__FILE__) %>"
                     assert(cast-><%= field.name %> != 0);
-                    argv[<%= index %>] = RARRAY_AREF(constants, cast-><%= field.name %> - 1);
+                    argv[<%= index %>] = RARRAY_AREF(constants, pm_constant_id_index(cast-><%= field.name %>));
                     <%- when Prism::Template::OptionalConstantField -%>
-                    argv[<%= index %>] = cast-><%= field.name %> == 0 ? Qnil : RARRAY_AREF(constants, cast-><%= field.name %> - 1);
+                    argv[<%= index %>] = cast-><%= field.name %> == 0 ? Qnil : RARRAY_AREF(constants, pm_constant_id_index(cast-><%= field.name %>));
                     <%- when Prism::Template::ConstantListField -%>
 #line <%= __LINE__ + 1 %> "prism/templates/ext/prism/<%= File.basename(__FILE__) %>"
                     argv[<%= index %>] = rb_ary_new_capa(cast-><%= field.name %>.size);
                     for (size_t index = 0; index < cast-><%= field.name %>.size; index++) {
                         assert(cast-><%= field.name %>.ids[index] != 0);
-                        rb_ary_push(argv[<%= index %>], RARRAY_AREF(constants, cast-><%= field.name %>.ids[index] - 1));
+                        rb_ary_push(argv[<%= index %>], RARRAY_AREF(constants, pm_constant_id_index(cast-><%= field.name %>.ids[index])));
                     }
                     if (freeze) rb_obj_freeze(argv[<%= index %>]);
                     <%- when Prism::Template::LocationField -%>