[PRISM] Don't calculate params size based on locals

Prior to this commit, we were conflating the size of the locals
list with the number of parameters. This commit distinguishes
the two, and fixes a related bug which would occur if we set a local
that was not a parameter

2 files changed, 19 insertions, 3 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 05c75ead50..1f16bfe044 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -710,13 +710,13 @@ pm_lookup_local_index(rb_iseq_t *iseq, pm_scope_node_t *scope_node, pm_constant_
 {
     st_data_t local_index;
 
-    int num_params = ISEQ_BODY(iseq)->param.size;
+    int locals_size = (int) scope_node->locals.size;
 
     if (!st_lookup(scope_node->index_lookup_table, constant_id, &local_index)) {
         rb_bug("This local does not exist");
     }
 
-    return num_params - (int)local_index;
+    return locals_size - (int)local_index;
 }
 
 static int
@@ -3592,7 +3592,6 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         VALUE idtmp = 0;
         rb_ast_id_table_t *tbl = ALLOCV(idtmp, sizeof(rb_ast_id_table_t) + locals_size * sizeof(ID));
         tbl->size = (int) locals_size;
-        body->param.size = (int) locals_size;
 
         for (size_t i = 0; i < locals_size; i++) {
             pm_constant_id_t constant_id = locals->ids[i];
@@ -3740,6 +3739,22 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             }
         }
 
+        iseq_calc_param_size(iseq);
+
+        // Calculating the parameter size above does not account for numbered parameters
+        // We can _only_ have numbered parameters if we don't have non numbered parameters
+        // We verify this through asserts, and add the numbered_parameters size accordingly
+        if (PM_NODE_TYPE_P(scope_node->ast_node, PM_BLOCK_NODE)) {
+            uint32_t numbered_parameters = ((pm_block_node_t *)(scope_node->ast_node))->numbered_parameters;
+            RUBY_ASSERT(!body->param.size || !numbered_parameters);
+            body->param.size += numbered_parameters;
+        }
+        else if (PM_NODE_TYPE_P(scope_node->ast_node, PM_LAMBDA_NODE)) {
+            uint32_t numbered_parameters = ((pm_lambda_node_t *)(scope_node->ast_node))->numbered_parameters;
+            RUBY_ASSERT(!body->param.size || !numbered_parameters);
+            body->param.size += numbered_parameters;
+        }
+
         iseq_set_local_table(iseq, tbl);
 
         switch (body->type) {
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index b7dc6a6e0f..8a078c01cc 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -756,6 +756,7 @@ module Prism
       assert_prism_eval("def self.prism_test_def_node() 1 end; prism_test_def_node()")
       assert_prism_eval("def self.prism_test_def_node(a,b) [a, b] end; prism_test_def_node(1,2)")
       assert_prism_eval("def self.prism_test_def_node(a,x=7,y=1) x end; prism_test_def_node(7,1)")
+      assert_prism_eval("def self.prism_test_def_node(a = 1); x = 2; end; prism_test_def_node")
 
       # rest argument
       assert_prism_eval("def self.prism_test_def_node(*a) a end; prism_test_def_node().inspect")