[PRISM] Restructure parameters

Prior to this commit, we weren't accounting for hidden variables
on the locals table, so we would have inconsistencies on the stack.
This commit fixes params, and introduces a hidden_variable_count
on the scope, both of which fix parameters.

3 files changed, 92 insertions, 15 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 9b24a33fee..1aca256f2e 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -704,7 +704,7 @@ pm_lookup_local_index_any_scope(rb_iseq_t *iseq, pm_scope_node_t *scope_node, pm
         return pm_lookup_local_index_any_scope(iseq, scope_node->previous, constant_id);
     }
 
-    return (int)scope_node->index_lookup_table->num_entries - (int)local_index;
+    return scope_node->hidden_variable_count + (int)scope_node->index_lookup_table->num_entries - (int)local_index;
 }
 
 static int
@@ -718,7 +718,7 @@ pm_lookup_local_index(rb_iseq_t *iseq, pm_scope_node_t *scope_node, pm_constant_
         rb_bug("This local does not exist");
     }
 
-    return locals_size - (int)local_index;
+    return scope_node->hidden_variable_count + locals_size - (int)local_index;
 }
 
 static int
@@ -1337,6 +1337,7 @@ pm_scope_node_init(const pm_node_t *node, pm_scope_node_t *scope, pm_scope_node_
     scope->body = NULL;
     scope->constants = NULL;
     scope->local_depth_offset = 0;
+    scope->hidden_variable_count = 0;
     if (previous) {
         scope->constants = previous->constants;
         scope->local_depth_offset = previous->local_depth_offset;
@@ -3721,9 +3722,16 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         // Index lookup table buffer size is only the number of the locals
         st_table *index_lookup_table = st_init_numtable();
 
+        int table_size = (int) locals_size;
+
+        if (keywords_list && keywords_list->size) {
+            table_size++;
+            scope_node->hidden_variable_count = 1;
+        }
+
         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;
+        rb_ast_id_table_t *tbl = ALLOCV(idtmp, sizeof(rb_ast_id_table_t) + table_size * sizeof(ID));
+        tbl->size = table_size;
 
         for (size_t i = 0; i < locals_size; i++) {
             pm_constant_id_t constant_id = locals->ids[i];
@@ -3738,10 +3746,19 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             st_insert(index_lookup_table, constant_id, i);
         }
 
+        // We have keywords and so we need to allocate
+        // space for another variable
+        if (table_size > (int) locals_size) {
+            tbl->ids[locals_size] = rb_make_temporary_id(locals_size);
+        }
+
         scope_node->index_lookup_table = index_lookup_table;
 
+        int arg_size = 0;
+
         if (optionals_list && optionals_list->size) {
             body->param.opt_num = (int) optionals_list->size;
+            arg_size += body->param.opt_num;
             LABEL **opt_table = (LABEL **)ALLOC_N(VALUE, optionals_list->size + 1);
             LABEL *label;
 
@@ -3768,6 +3785,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
 
         if (requireds_list && requireds_list->size) {
             body->param.lead_num = (int) requireds_list->size;
+            arg_size += body->param.lead_num;
             body->param.flags.has_lead = true;
 
             for (size_t i = 0; i < requireds_list->size; i++) {
@@ -3779,16 +3797,25 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             }
         }
 
+        if (parameters_node && parameters_node->rest) {
+            body->param.rest_start = arg_size++;
+            body->param.flags.has_rest = true;
+            assert(body->param.rest_start != -1);
+        }
+
         if (posts_list && posts_list->size) {
             body->param.post_num = (int) posts_list->size;
-            body->param.post_start = body->param.lead_num + body->param.opt_num; // TODO
+            body->param.post_start = arg_size;
             body->param.flags.has_post = true;
+            arg_size += body->param.post_num;
         }
 
         // Keywords create an internal variable on the parse tree
         if (keywords_list && keywords_list->size) {
             body->param.keyword = keyword = ZALLOC_N(struct rb_iseq_param_keyword, 1);
             keyword->num = (int) keywords_list->size;
+            arg_size += keyword->num;
+            keyword->bits_start = arg_size++;
 
             body->param.flags.has_kw = true;
             const VALUE default_values = rb_ary_hidden_new(1);
@@ -3811,7 +3838,18 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
                           rb_ary_push(default_values, pm_static_literal_value(value, scope_node, parser));
                       }
                       else {
-                          PM_COMPILE_POPPED(value);
+                          LABEL *end_label = NEW_LABEL(nd_line(&dummy_line_node));
+
+                          int index = pm_lookup_local_index(iseq, scope_node, name);
+                          int kw_bits_idx = table_size - body->param.keyword->bits_start;
+                          int keyword_idx = (int)(i + scope_node->hidden_variable_count);
+
+                          ADD_INSN2(ret, &dummy_line_node, checkkeyword, INT2FIX(kw_bits_idx + VM_ENV_DATA_SIZE - 1), INT2FIX(keyword_idx - 1));
+                          ADD_INSNL(ret, &dummy_line_node, branchif, end_label);
+                          PM_COMPILE(value);
+                          ADD_SETLOCAL(ret, &dummy_line_node, index, 0);
+
+                          ADD_LABEL(ret, end_label);
                           rb_ary_push(default_values, complex_mark);
                       }
 
@@ -3847,31 +3885,31 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         }
 
         if (parameters_node) {
-            if (parameters_node->rest) {
-                body->param.rest_start = body->param.lead_num + body->param.opt_num;
-                body->param.flags.has_rest = true;
-            }
-
             if (parameters_node->keyword_rest) {
                 if (PM_NODE_TYPE_P(parameters_node->keyword_rest, PM_NO_KEYWORDS_PARAMETER_NODE)) {
                     body->param.flags.accepts_no_kwarg = true;
                 }
                 else {
-                    if (!body->param.flags.has_kw) {
+                    if (body->param.flags.has_kw) {
+                        arg_size--;
+                    }
+                    else {
                         body->param.keyword = keyword = ZALLOC_N(struct rb_iseq_param_keyword, 1);
                     }
-                    keyword->rest_start = (int) locals_size - (parameters_node->block ? 2 : 1);
+
+                    keyword->rest_start = arg_size++;
                     body->param.flags.has_kwrest = true;
                 }
             }
 
             if (parameters_node->block) {
-                body->param.block_start = (int) locals_size - 1;
+                body->param.block_start = arg_size;//(int) locals_size - 1;
                 body->param.flags.has_block = true;
             }
         }
 
         iseq_calc_param_size(iseq);
+        body->param.size = arg_size;
 
         // Calculating the parameter size above does not account for numbered
         // parameters. We can _only_ have numbered parameters if we don't have
diff --git a/prism_compile.h b/prism_compile.h
index 2d64811e69..7f45ddd1d2 100644
--- a/prism_compile.h
+++ b/prism_compile.h
@@ -11,6 +11,14 @@ typedef struct pm_scope_node {
     pm_constant_id_list_t locals;
     pm_parser_t *parser;
 
+    // There are sometimes when we need to track
+    // hidden variables that we have put on
+    // the local table for the stack to use, so
+    // that we properly account for them when giving
+    // local indexes. We do this with the
+    // hidden_variable_count
+    int hidden_variable_count;
+
     ID *constants;
     st_table *index_lookup_table;
 
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index d62ee7c879..c56792de0d 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -879,7 +879,17 @@ module Prism
       assert_prism_eval("def self.prism_test_def_node(x,y,z=7,*a) a end; prism_test_def_node(7,7).inspect")
       assert_prism_eval("def self.prism_test_def_node(x,y,z=7,zz=7,*a) a end; prism_test_def_node(7,7,7).inspect")
 
-      # block argument
+      # keyword arguments
+      assert_prism_eval("def self.prism_test_def_node(a: 1, b: 2, c: 4) a + b + c; end; prism_test_def_node(a: 2)")
+      assert_prism_eval("def self.prism_test_def_node(a: 1, b: 2, c: 4) a + b + c; end; prism_test_def_node(b: 3)")
+      assert_prism_eval(<<-CODE)
+        def self.prism_test_def_node(x = 1, y, a: 8, b: 2, c: 4)
+          a + b + c + x + y
+        end
+        prism_test_def_node(10, b: 3)
+      CODE
+
+      # block arguments
       assert_prism_eval("def self.prism_test_def_node(&block) block end; prism_test_def_node{}.class")
       assert_prism_eval("def self.prism_test_def_node(&block) block end; prism_test_def_node().inspect")
       assert_prism_eval("def self.prism_test_def_node(a,b=7,*c,&block) b end; prism_test_def_node(7,1).inspect")
@@ -914,6 +924,27 @@ module Prism
 
         method(:prism_test_method_parameters).parameters
       CODE
+
+      assert_prism_eval(<<-CODE)
+        def self.prism_test_method_parameters(d:, e: 2, **f, &g)
+        end
+
+        method(:prism_test_method_parameters).parameters
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        def self.prism_test_method_parameters(**f, &g)
+        end
+
+        method(:prism_test_method_parameters).parameters
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        def self.prism_test_method_parameters(&g)
+        end
+
+        method(:prism_test_method_parameters).parameters
+      CODE
     end
 
     def test_LambdaNode