Merge pull request #2357 from ruby/constant-pool-find

Provide APIs for finding value in constant pool

include/prism/util/pm_constant_pool.h

@@ -154,6 +154,17 @@ bool pm_constant_pool_init(pm_constant_pool_t *pool, uint32_t capacity);
  */
 pm_constant_t * pm_constant_pool_id_to_constant(const pm_constant_pool_t *pool, pm_constant_id_t constant_id);
 
+/**
+ * Find a constant in a constant pool. Returns the id of the constant, or 0 if
+ * the constant is not found.
+ *
+ * @param pool The pool to find the constant in.
+ * @param start A pointer to the start of the constant.
+ * @param length The length of the constant.
+ * @return The id of the constant.
+ */
+pm_constant_id_t pm_constant_pool_find(pm_constant_pool_t *pool, const uint8_t *start, size_t length);
+
 /**
  * Insert a constant into a constant pool that is a slice of a source string.
  * Returns the id of the constant, or 0 if any potential calls to resize fail.

src/prism.c

@@ -6325,6 +6325,16 @@ pm_parser_local_add_owned(pm_parser_t *parser, const uint8_t *start, size_t leng
     return constant_id;
 }
 
+/**
+ * Add a local variable from a constant string to the current scope.
+ */
+static pm_constant_id_t
+pm_parser_local_add_constant(pm_parser_t *parser, const char *start, size_t length) {
+    pm_constant_id_t constant_id = pm_parser_constant_id_constant(parser, start, length);
+    if (constant_id != 0) pm_parser_local_add(parser, constant_id);
+    return constant_id;
+}
+
 /**
  * Add a parameter name to the current scope and check whether the name of the
  * parameter is unique or not.
@@ -13195,6 +13205,15 @@ outer_scope_using_numbered_parameters_p(pm_parser_t *parser) {
     return false;
 }
 
+/**
+ * These are the names of the various numbered parameters. We have them here so
+ * that when we insert them into the constant pool we can use a constant string
+ * and not have to allocate.
+ */
+static const char * const pm_numbered_parameter_names[] = {
+    "_1", "_2", "_3", "_4", "_5", "_6", "_7", "_8", "_9"
+};
+
 /**
  * Parse an identifier into either a local variable read. If the local variable
  * is not found, it returns NULL instead.
@@ -13217,12 +13236,10 @@ parse_variable(pm_parser_t *parser) {
             pm_parser_err_previous(parser, PM_ERR_NUMBERED_PARAMETER_OUTER_SCOPE);
         } else {
             // Indicate that this scope is using numbered params so that child
-            // scopes cannot.
-            uint8_t number = parser->previous.start[1];
-
-            // We subtract the value for the character '0' to get the actual
-            // integer value of the number (only _1 through _9 are valid)
-            uint8_t numbered_parameters = (uint8_t) (number - '0');
+            // scopes cannot. We subtract the value for the character '0' to get
+            // the actual integer value of the number (only _1 through _9 are
+            // valid).
+            uint8_t numbered_parameters = (uint8_t) (parser->previous.start[1] - '0');
             if (numbered_parameters > parser->current_scope->numbered_parameters) {
                 parser->current_scope->numbered_parameters = numbered_parameters;
                 pm_parser_numbered_parameters_set(parser, numbered_parameters);
@@ -13233,21 +13250,13 @@ parse_variable(pm_parser_t *parser) {
             // referencing _2 means that _1 must exist. Therefore here we
             // loop through all of the possibilities and add them into the
             // constant pool.
-            uint8_t current = '1';
-            uint8_t *value;
-
-            while (current < number) {
-                value = malloc(2);
-                value[0] = '_';
-                value[1] = current++;
-                pm_parser_local_add_owned(parser, value, 2);
+            for (uint8_t numbered_parameter = 1; numbered_parameter <= numbered_parameters - 1; numbered_parameter++) {
+                pm_parser_local_add_constant(parser, pm_numbered_parameter_names[numbered_parameter - 1], 2);
             }
 
-            // Now we can add the actual token that is being used. For
-            // this one we can add a shared version since it is directly
-            // referenced in the source.
-            pm_parser_local_add_token(parser, &parser->previous);
-            return pm_local_variable_read_node_create(parser, &parser->previous, 0);
+            // Finally we can create the local variable read node.
+            pm_constant_id_t name_id = pm_parser_local_add_constant(parser, pm_numbered_parameter_names[numbered_parameters - 1], 2);
+            return pm_local_variable_read_node_create_constant_id(parser, &parser->previous, name_id, 0);
         }
     }
 

src/util/pm_constant_pool.c

@@ -181,6 +181,31 @@ pm_constant_pool_id_to_constant(const pm_constant_pool_t *pool, pm_constant_id_t
     return &pool->constants[constant_id - 1];
 }
 
+/**
+ * Find a constant in a constant pool. Returns the id of the constant, or 0 if
+ * the constant is not found.
+ */
+pm_constant_id_t
+pm_constant_pool_find(pm_constant_pool_t *pool, const uint8_t *start, size_t length) {
+    assert(is_power_of_two(pool->capacity));
+    const uint32_t mask = pool->capacity - 1;
+
+    uint32_t hash = pm_constant_pool_hash(start, length);
+    uint32_t index = hash & mask;
+    pm_constant_pool_bucket_t *bucket;
+
+    while (bucket = &pool->buckets[index], bucket->id != PM_CONSTANT_ID_UNSET) {
+        pm_constant_t *constant = &pool->constants[bucket->id - 1];
+        if ((constant->length == length) && memcmp(constant->start, start, length) == 0) {
+            return bucket->id;
+        }
+
+        index = (index + 1) & mask;
+    }
+
+    return PM_CONSTANT_ID_UNSET;
+}
+
 /**
  * Insert a constant into a constant pool and return its index in the pool.
  */