Move logic to emit_undo_profile_instruction_count

solana-labs · Nov 18, 2024 · deec3ad · deec3ad
1 parent 08cac7f
commit deec3ad
Showing 1 changed file with 19 additions and 12 deletions.
diff --git a/src/jit.rs b/src/jit.rs
@@ -979,9 +979,21 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
     }
 
     #[inline]
-    fn emit_undo_profile_instruction_count(&mut self, target_pc: usize) {
+    fn emit_undo_profile_instruction_count(&mut self, target_pc: Option<usize>) {
         if self.config.enable_instruction_meter {
-            self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, self.pc as i64 + 1 - target_pc as i64, None)); // instruction_meter += (self.pc + 1) - target_pc;
+            match target_pc {
+                Some(target_pc) => {
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, self.pc as i64 + 1 - target_pc as i64, None)); // instruction_meter += (self.pc + 1) - target_pc;
+                }
+                None => {
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x29, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter -= guest_target_pc
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, 1, None)); // instruction_meter += 1
+                    // Retrieve the current program from the stack. `return_near` popped an element from the stack,
+                    // so the offset is 16. Check `ANCHOR_INTERNAL_FUNCTION_CALL_REG` for more details.
+                    self.emit_ins(X86Instruction::load(OperandSize::S64, RSP, REGISTER_SCRATCH, X86IndirectAccess::OffsetIndexShift(-16, RSP, 0)));
+                    self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x01, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter += guest_current_pc
+                }
+            }
         }
     }
 
@@ -1125,7 +1137,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
             }
         }
 
-        self.emit_undo_profile_instruction_count(0);
+        self.emit_undo_profile_instruction_count(Some(0));
 
         // Restore the previous frame pointer
         self.emit_ins(X86Instruction::pop(REGISTER_MAP[FRAME_PTR_REG]));
@@ -1139,7 +1151,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_validate_and_profile_instruction_count(false, Some(0));
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, function as usize as i64));
         self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL, 5)));
-        self.emit_undo_profile_instruction_count(0);
+        self.emit_undo_profile_instruction_count(Some(0));
     }
 
     #[inline]
@@ -1224,7 +1236,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, target_pc as i64));
         let jump_offset = self.relative_to_target_pc(target_pc, 6);
         self.emit_ins(X86Instruction::conditional_jump_immediate(op, jump_offset));
-        self.emit_undo_profile_instruction_count(target_pc);
+        self.emit_undo_profile_instruction_count(Some(target_pc));
     }
 
     #[inline]
@@ -1245,7 +1257,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, target_pc as i64));
         let jump_offset = self.relative_to_target_pc(target_pc, 6);
         self.emit_ins(X86Instruction::conditional_jump_immediate(op, jump_offset));
-        self.emit_undo_profile_instruction_count(target_pc);
+        self.emit_undo_profile_instruction_count(Some(target_pc));
     }
 
     fn emit_shift(&mut self, size: OperandSize, opcode_extension: u8, source: u8, destination: u8, immediate: Option<i64>) {
@@ -1590,12 +1602,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
         // If callx lands in an invalid address, we must undo the changes in the instruction meter
         // so that we can correctly calculate the number of executed instructions for error handling.
         self.set_anchor(ANCHOR_CALL_REG_UNSUPPORTED_INSTRUCTION);
-        self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x29, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter -= guest_target_pc
-        self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, REGISTER_INSTRUCTION_METER, 1, None)); // instruction_meter += 1
-        // Retrieve the current program from the stack. `return_near` popped an element from the stack,
-        // so the offset is 16.
-        self.emit_ins(X86Instruction::load(OperandSize::S64, RSP, REGISTER_SCRATCH, X86IndirectAccess::OffsetIndexShift(-16, RSP, 0)));
-        self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x01, REGISTER_SCRATCH, REGISTER_INSTRUCTION_METER, 0, None)); // instruction_meter += guest_current_pc
+        self.emit_undo_profile_instruction_count(None);
         self.emit_ins(X86Instruction::jump_immediate(self.relative_to_anchor(ANCHOR_CALL_UNSUPPORTED_INSTRUCTION, 5)));
 
         // Translates a vm memory address to a host memory address