RJIT: Add --rjit-verify-ctx option

6 files changed, 99 insertions, 16 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index 203c2658df..1c024c4c40 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -279,6 +279,8 @@ module RubyVM::RJIT
     end
 
     # @param asm [RubyVM::RJIT::Assembler]
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
     def compile_block(asm, jit:, pc:, ctx: Context.new)
       # Mark the block start address and prepare an exit code storage
       ctx = limit_block_versions(jit.iseq, pc, ctx)
@@ -292,6 +294,7 @@ module RubyVM::RJIT
       # Compile each insn
       index = (pc - iseq.body.iseq_encoded.to_i) / C.VALUE.size
       while index < iseq.body.iseq_size
+        # Set the current instruction
         insn = self.class.decode_insn(iseq.body.iseq_encoded[index])
         jit.pc = (iseq.body.iseq_encoded + index).to_i
         jit.stack_size_for_pc = ctx.stack_size
@@ -308,6 +311,11 @@ module RubyVM::RJIT
           jit.record_boundary_patch_point = false
         end
 
+        # In debug mode, verify our existing assumption
+        if C.rjit_opts.verify_ctx && jit.at_current_insn?
+          verify_ctx(jit, ctx)
+        end
+
         case status = @insn_compiler.compile(jit, ctx, asm, insn)
         when KeepCompiling
           # For now, reset the chain depth after each instruction as only the
@@ -435,5 +443,67 @@ module RubyVM::RJIT
     rescue RangeError # bignum too big to convert into `unsigned long long' (RangeError)
       -1
     end
+
+    # Verify the ctx's types and mappings against the compile-time stack, self, and locals.
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    def verify_ctx(jit, ctx)
+      # Only able to check types when at current insn
+      assert(jit.at_current_insn?)
+
+      self_val = jit.peek_at_self
+      self_val_type = Type.from(self_val)
+
+      # Verify self operand type
+      assert_compatible(self_val_type, ctx.get_opnd_type(SelfOpnd))
+
+      # Verify stack operand types
+      [ctx.stack_size, MAX_TEMP_TYPES].min.times do |i|
+        learned_mapping, learned_type = ctx.get_opnd_mapping(StackOpnd[i])
+        stack_val = jit.peek_at_stack(i)
+        val_type = Type.from(stack_val)
+
+        case learned_mapping
+        in MapToSelf
+          if C.to_value(self_val) != C.to_value(stack_val)
+            raise "verify_ctx: stack value was mapped to self, but values did not match:\n"\
+              "stack: #{stack_val.inspect}, self: #{self_val.inspect}"
+          end
+        in MapToLocal[local_idx]
+          local_val = jit.peek_at_local(local_idx)
+          if C.to_value(local_val) != C.to_value(stack_val)
+            raise "verify_ctx: stack value was mapped to local, but values did not match:\n"\
+              "stack: #{stack_val.inspect}, local: #{local_val.inspect}"
+          end
+        in MapToStack
+          # noop
+        end
+
+        # If the actual type differs from the learned type
+        assert_compatible(val_type, learned_type)
+      end
+
+      # Verify local variable types
+      local_table_size = jit.iseq.body.local_table_size
+      [local_table_size, MAX_TEMP_TYPES].min.times do |i|
+        learned_type = ctx.get_local_type(i)
+        local_val = jit.peek_at_local(i)
+        local_type = Type.from(local_val)
+
+        assert_compatible(local_type, learned_type)
+      end
+    end
+
+    def assert_compatible(actual_type, ctx_type)
+      if actual_type.diff(ctx_type) == TypeDiff::Incompatible
+        raise "verify_ctx: ctx type (#{ctx_type.type.inspect}) is incompatible with actual type (#{actual_type.type.inspect})"
+      end
+    end
+
+    def assert(cond)
+      unless cond
+        raise "'#{cond.inspect}' was not true"
+      end
+    end
   end
 end
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index e3d3c36836..d9aca5b52c 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -21,16 +21,16 @@ module RubyVM::RJIT
     # @param insn `RubyVM::RJIT::Instruction`
     def compile(jit, ctx, asm, insn)
       asm.incr_counter(:rjit_insns_count)
-      insn_idx = format('%04d', (jit.pc.to_i - jit.iseq.body.iseq_encoded.to_i) / C.VALUE.size)
-      asm.comment("Insn: #{insn_idx} #{insn.name}")
 
-      # stack = ctx.stack_size.times.map do |stack_idx|
-      #   ctx.get_opnd_type(StackOpnd[ctx.stack_size - stack_idx - 1]).type
-      # end
-      # locals = jit.iseq.body.local_table_size.times.map do |local_idx|
-      #   (ctx.local_types[local_idx] || Type::Unknown).type
-      # end
-      # asm.comment("Insn: #{insn_idx} #{insn.name} (stack: [#{stack.join(', ')}], locals: [#{locals.join(', ')}])")
+      stack = ctx.stack_size.times.map do |stack_idx|
+        ctx.get_opnd_type(StackOpnd[ctx.stack_size - stack_idx - 1]).type
+      end
+      locals = jit.iseq.body.local_table_size.times.map do |local_idx|
+        (ctx.local_types[local_idx] || Type::Unknown).type
+      end
+
+      insn_idx = format('%04d', (jit.pc.to_i - jit.iseq.body.iseq_encoded.to_i) / C.VALUE.size)
+      asm.comment("Insn: #{insn_idx} #{insn.name} (stack: [#{stack.join(', ')}], locals: [#{locals.join(', ')}])")
 
       # 83/102
       case insn.name
diff --git a/lib/ruby_vm/rjit/jit_state.rb b/lib/ruby_vm/rjit/jit_state.rb
index f31223e478..02a713474e 100644
--- a/lib/ruby_vm/rjit/jit_state.rb
+++ b/lib/ruby_vm/rjit/jit_state.rb
@@ -27,6 +27,13 @@ module RubyVM::RJIT
       pc == cfp.pc.to_i
     end
 
+    def peek_at_local(n)
+      local_table_size = iseq.body.local_table_size
+      offset = -C::VM_ENV_DATA_SIZE - local_table_size + n + 1
+      value = (cfp.ep + offset).*
+      C.to_ruby(value)
+    end
+
     def peek_at_stack(depth_from_top)
       raise 'not at current insn' unless at_current_insn?
       offset = -(1 + depth_from_top)
diff --git a/rjit.c b/rjit.c
index f1d3440ce9..e17e4b982e 100644
--- a/rjit.c
+++ b/rjit.c
@@ -118,25 +118,28 @@ rb_rjit_setup_options(const char *s, struct rb_rjit_options *rjit_opt)
     if (l == 0) {
         return;
     }
+    else if (opt_match_arg(s, l, "call-threshold")) {
+        rjit_opt->call_threshold = atoi(s + 1);
+    }
+    else if (opt_match_arg(s, l, "exec-mem-size")) {
+        rjit_opt->exec_mem_size = atoi(s + 1);
+    }
     else if (opt_match_noarg(s, l, "stats")) {
         rjit_opt->stats = true;
     }
     else if (opt_match_noarg(s, l, "trace-exits")) {
         rjit_opt->trace_exits = true;
     }
-    else if (opt_match_arg(s, l, "call-threshold")) {
-        rjit_opt->call_threshold = atoi(s + 1);
+    else if (opt_match_noarg(s, l, "dump-disasm")) {
+        rjit_opt->dump_disasm = true;
     }
-    else if (opt_match_arg(s, l, "exec-mem-size")) {
-        rjit_opt->exec_mem_size = atoi(s + 1);
+    else if (opt_match_noarg(s, l, "verify-ctx")) {
+        rjit_opt->verify_ctx = true;
     }
     // --rjit=pause is an undocumented feature for experiments
     else if (opt_match_noarg(s, l, "pause")) {
         rjit_opt->pause = true;
     }
-    else if (opt_match_noarg(s, l, "dump-disasm")) {
-        rjit_opt->dump_disasm = true;
-    }
     else {
         rb_raise(rb_eRuntimeError,
                  "invalid RJIT option `%s' (--help will show valid RJIT options)", s);
diff --git a/rjit.h b/rjit.h
index aeb0a933df..1efd4f3ae7 100644
--- a/rjit.h
+++ b/rjit.h
@@ -36,6 +36,8 @@ struct rb_rjit_options {
     bool trace_exits;
     // Enable disasm of all JIT code
     bool dump_disasm;
+    // Verify context objects
+    bool verify_ctx;
     // [experimental] Do not start RJIT until RJIT.resume is called.
     bool pause;
 };
diff --git a/rjit_c.rb b/rjit_c.rb
index fa4ea2dddc..843ee53295 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -1316,6 +1316,7 @@ module RubyVM::RJIT # :nodoc: all
       stats: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_options *)NULL)), stats)")],
       trace_exits: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_options *)NULL)), trace_exits)")],
       dump_disasm: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_options *)NULL)), dump_disasm)")],
+      verify_ctx: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_options *)NULL)), verify_ctx)")],
       pause: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_options *)NULL)), pause)")],
     )
   end