Merge pull request #203 from jhawthorn/getspecial

maximecb · web-flow · commit 63e927534527 · 2021-09-13T11:40:50.000-07:00
Implement getspecial
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
@@ -205,6 +205,15 @@ def test_expandarray_nil
     RUBY
   end
 
+  def test_getspecial_backref
+    assert_compiles("'foo' =~ /(o)./; $&", insns: %i[getspecial], result: "oo")
+    assert_compiles("'foo' =~ /(o)./; $`", insns: %i[getspecial], result: "f")
+    assert_compiles("'foo' =~ /(o)./; $'", insns: %i[getspecial], result: "")
+    assert_compiles("'foo' =~ /(o)./; $+", insns: %i[getspecial], result: "o")
+    assert_compiles("'foo' =~ /(o)./; $1", insns: %i[getspecial], result: "o")
+    assert_compiles("'foo' =~ /(o)./; $2", insns: %i[getspecial], result: nil)
+  end
+
   def test_compile_opt_getinlinecache
     assert_compiles(<<~RUBY, insns: %i[opt_getinlinecache], result: 123, min_calls: 2)
       def get_foo
diff --git a/yjit_codegen.c b/yjit_codegen.c
@@ -3806,6 +3806,77 @@ gen_toregexp(jitstate_t* jit, ctx_t* ctx)
     return YJIT_KEEP_COMPILING;
 }
 
+static codegen_status_t
+gen_getspecial(jitstate_t *jit, ctx_t *ctx)
+{
+    // This takes two arguments, key and type
+    // key is only used when type == 0
+    // A non-zero type determines which type of backref to fetch
+    rb_num_t key = jit_get_arg(jit, 0);
+    rb_num_t type = jit_get_arg(jit, 1);
+
+    if (type == 0) {
+        // not yet implemented
+        return YJIT_CANT_COMPILE;
+    } else if (type & 0x01) {
+        // Fetch a "special" backref based on a char encoded by shifting by 1
+
+        // Can raise if matchdata uninitialized
+        jit_prepare_routine_call(jit, ctx, REG0);
+
+        // call rb_backref_get()
+        ADD_COMMENT(cb, "rb_backref_get");
+        call_ptr(cb, REG0, (void *)rb_backref_get);
+        mov(cb, C_ARG_REGS[0], RAX);
+
+        switch (type >> 1) {
+            case '&':
+                ADD_COMMENT(cb, "rb_reg_last_match");
+                call_ptr(cb, REG0, (void *)rb_reg_last_match);
+                break;
+            case '`':
+                ADD_COMMENT(cb, "rb_reg_match_pre");
+                call_ptr(cb, REG0, (void *)rb_reg_match_pre);
+                break;
+            case '\'':
+                ADD_COMMENT(cb, "rb_reg_match_post");
+                call_ptr(cb, REG0, (void *)rb_reg_match_post);
+                break;
+            case '+':
+                ADD_COMMENT(cb, "rb_reg_match_last");
+                call_ptr(cb, REG0, (void *)rb_reg_match_last);
+                break;
+            default:
+                rb_bug("invalid back-ref");
+        }
+
+        x86opnd_t stack_ret = ctx_stack_push(ctx, TYPE_UNKNOWN);
+        mov(cb, stack_ret, RAX);
+
+        return YJIT_KEEP_COMPILING;
+    } else {
+        // Fetch the N-th match from the last backref based on type shifted by 1
+
+        // Can raise if matchdata uninitialized
+        jit_prepare_routine_call(jit, ctx, REG0);
+
+        // call rb_backref_get()
+        ADD_COMMENT(cb, "rb_backref_get");
+        call_ptr(cb, REG0, (void *)rb_backref_get);
+
+        // rb_reg_nth_match((int)(type >> 1), backref);
+        ADD_COMMENT(cb, "rb_reg_nth_match");
+        mov(cb, C_ARG_REGS[0], imm_opnd(type >> 1));
+        mov(cb, C_ARG_REGS[1], RAX);
+        call_ptr(cb, REG0, (void *)rb_reg_nth_match);
+
+        x86opnd_t stack_ret = ctx_stack_push(ctx, TYPE_UNKNOWN);
+        mov(cb, stack_ret, RAX);
+
+        return YJIT_KEEP_COMPILING;
+    }
+}
+
 static codegen_status_t
 gen_opt_getinlinecache(jitstate_t *jit, ctx_t *ctx)
 {
@@ -4156,6 +4227,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(setglobal), gen_setglobal);
     yjit_reg_op(BIN(tostring), gen_tostring);
     yjit_reg_op(BIN(toregexp), gen_toregexp);
+    yjit_reg_op(BIN(getspecial), gen_getspecial);
 
     yjit_method_codegen_table = st_init_numtable();
 
