Unify BINARY_OP bytecodes (#6317)

* Unify BINARY_OP bytecodes

* Add missing op to `as_inplace`

* Fix doc example

* Fix jit

* Fix doc

* Use correct opname

* Fix dis fmt

* Inplace ops support in JIT

crates/codegen/src/compile.rs

@@ -3301,7 +3301,7 @@ impl Compiler {
                 // Subtract to compute the correct index.
                 emit!(
                     self,
-                    Instruction::BinaryOperation {
+                    Instruction::BinaryOp {
                         op: BinaryOperator::Subtract
                     }
                 );
@@ -4339,26 +4339,24 @@ impl Compiler {
     }
 
     fn compile_op(&mut self, op: &Operator, inplace: bool) {
-        let op = match op {
-            Operator::Add => bytecode::BinaryOperator::Add,
-            Operator::Sub => bytecode::BinaryOperator::Subtract,
-            Operator::Mult => bytecode::BinaryOperator::Multiply,
-            Operator::MatMult => bytecode::BinaryOperator::MatrixMultiply,
-            Operator::Div => bytecode::BinaryOperator::Divide,
-            Operator::FloorDiv => bytecode::BinaryOperator::FloorDivide,
-            Operator::Mod => bytecode::BinaryOperator::Modulo,
-            Operator::Pow => bytecode::BinaryOperator::Power,
-            Operator::LShift => bytecode::BinaryOperator::Lshift,
-            Operator::RShift => bytecode::BinaryOperator::Rshift,
-            Operator::BitOr => bytecode::BinaryOperator::Or,
-            Operator::BitXor => bytecode::BinaryOperator::Xor,
-            Operator::BitAnd => bytecode::BinaryOperator::And,
+        let bin_op = match op {
+            Operator::Add => BinaryOperator::Add,
+            Operator::Sub => BinaryOperator::Subtract,
+            Operator::Mult => BinaryOperator::Multiply,
+            Operator::MatMult => BinaryOperator::MatrixMultiply,
+            Operator::Div => BinaryOperator::TrueDivide,
+            Operator::FloorDiv => BinaryOperator::FloorDivide,
+            Operator::Mod => BinaryOperator::Remainder,
+            Operator::Pow => BinaryOperator::Power,
+            Operator::LShift => BinaryOperator::Lshift,
+            Operator::RShift => BinaryOperator::Rshift,
+            Operator::BitOr => BinaryOperator::Or,
+            Operator::BitXor => BinaryOperator::Xor,
+            Operator::BitAnd => BinaryOperator::And,
         };
-        if inplace {
-            emit!(self, Instruction::BinaryOperationInplace { op })
-        } else {
-            emit!(self, Instruction::BinaryOperation { op })
-        }
+
+        let op = if inplace { bin_op.as_inplace() } else { bin_op };
+        emit!(self, Instruction::BinaryOp { op })
     }
 
     /// Implement boolean short circuit evaluation logic.

crates/compiler-core/src/bytecode.rs

@@ -594,10 +594,7 @@ pub enum Instruction {
     UnaryOperation {
         op: Arg<UnaryOperator>,
     },
-    BinaryOperation {
-        op: Arg<BinaryOperator>,
-    },
-    BinaryOperationInplace {
+    BinaryOp {
         op: Arg<BinaryOperator>,
     },
     BinarySubscript,
@@ -1094,32 +1091,142 @@ op_arg_enum!(
 
 op_arg_enum!(
     /// The possible Binary operators
+    ///
     /// # Examples
     ///
-    /// ```ignore
-    /// use rustpython_compiler_core::Instruction::BinaryOperation;
-    /// use rustpython_compiler_core::BinaryOperator::Add;
-    /// let op = BinaryOperation {op: Add};
+    /// ```rust
+    /// use rustpython_compiler_core::bytecode::{Arg, BinaryOperator, Instruction};
+    /// let (op, _) = Arg::new(BinaryOperator::Add);
+    /// let instruction = Instruction::BinaryOp { op };
     /// ```
-    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
+    ///
+    /// See also:
+    /// - [_PyEval_BinaryOps](https://github.com/python/cpython/blob/8183fa5e3f78ca6ab862de7fb8b14f3d929421e0/Python/ceval.c#L316-L343)
     #[repr(u8)]
+    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
     pub enum BinaryOperator {
-        Power = 0,
-        Multiply = 1,
-        MatrixMultiply = 2,
-        Divide = 3,
-        FloorDivide = 4,
-        Modulo = 5,
-        Add = 6,
-        Subtract = 7,
-        Lshift = 8,
+        /// `+`
+        Add = 0,
+        /// `&`
+        And = 1,
+        /// `//`
+        FloorDivide = 2,
+        /// `<<`
+        Lshift = 3,
+        /// `@`
+        MatrixMultiply = 4,
+        /// `*`
+        Multiply = 5,
+        /// `%`
+        Remainder = 6,
+        /// `|`
+        Or = 7,
+        /// `**`
+        Power = 8,
+        /// `>>`
         Rshift = 9,
-        And = 10,
-        Xor = 11,
-        Or = 12,
+        /// `-`
+        Subtract = 10,
+        /// `/`
+        TrueDivide = 11,
+        /// `^`
+        Xor = 12,
+        /// `+=`
+        InplaceAdd = 13,
+        /// `&=`
+        InplaceAnd = 14,
+        /// `//=`
+        InplaceFloorDivide = 15,
+        /// `<<=`
+        InplaceLshift = 16,
+        /// `@=`
+        InplaceMatrixMultiply = 17,
+        /// `*=`
+        InplaceMultiply = 18,
+        /// `%=`
+        InplaceRemainder = 19,
+        /// `|=`
+        InplaceOr = 20,
+        /// `**=`
+        InplacePower = 21,
+        /// `>>=`
+        InplaceRshift = 22,
+        /// `-=`
+        InplaceSubtract = 23,
+        /// `/=`
+        InplaceTrueDivide = 24,
+        /// `^=`
+        InplaceXor = 25,
     }
 );
 
+impl BinaryOperator {
+    /// Get the "inplace" version of the operator.
+    /// This has no effect if `self` is already an "inplace" operator.
+    ///
+    /// # Example
+    /// ```rust
+    /// use rustpython_compiler_core::bytecode::BinaryOperator;
+    ///
+    /// assert_eq!(BinaryOperator::Power.as_inplace(), BinaryOperator::InplacePower);
+    ///
+    /// assert_eq!(BinaryOperator::InplaceSubtract.as_inplace(), BinaryOperator::InplaceSubtract);
+    /// ```
+    #[must_use]
+    pub const fn as_inplace(self) -> Self {
+        match self {
+            Self::Add => Self::InplaceAdd,
+            Self::And => Self::InplaceAnd,
+            Self::FloorDivide => Self::InplaceFloorDivide,
+            Self::Lshift => Self::InplaceLshift,
+            Self::MatrixMultiply => Self::InplaceMatrixMultiply,
+            Self::Multiply => Self::InplaceMultiply,
+            Self::Remainder => Self::InplaceRemainder,
+            Self::Or => Self::InplaceOr,
+            Self::Power => Self::InplacePower,
+            Self::Rshift => Self::InplaceRshift,
+            Self::Subtract => Self::InplaceSubtract,
+            Self::TrueDivide => Self::InplaceTrueDivide,
+            Self::Xor => Self::InplaceXor,
+            _ => self,
+        }
+    }
+}
+
+impl fmt::Display for BinaryOperator {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let op = match self {
+            Self::Add => "+",
+            Self::And => "&",
+            Self::FloorDivide => "//",
+            Self::Lshift => "<<",
+            Self::MatrixMultiply => "@",
+            Self::Multiply => "*",
+            Self::Remainder => "%",
+            Self::Or => "|",
+            Self::Power => "**",
+            Self::Rshift => ">>",
+            Self::Subtract => "-",
+            Self::TrueDivide => "/",
+            Self::Xor => "^",
+            Self::InplaceAdd => "+=",
+            Self::InplaceAnd => "&=",
+            Self::InplaceFloorDivide => "//=",
+            Self::InplaceLshift => "<<=",
+            Self::InplaceMatrixMultiply => "@=",
+            Self::InplaceMultiply => "*=",
+            Self::InplaceRemainder => "%=",
+            Self::InplaceOr => "|=",
+            Self::InplacePower => "**=",
+            Self::InplaceRshift => ">>=",
+            Self::InplaceSubtract => "-=",
+            Self::InplaceTrueDivide => "/=",
+            Self::InplaceXor => "^=",
+        };
+        write!(f, "{op}")
+    }
+}
+
 op_arg_enum!(
     /// The possible unary operators
     #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -1514,7 +1621,7 @@ impl Instruction {
             DeleteAttr { .. } => -1,
             LoadConst { .. } => 1,
             UnaryOperation { .. } => 0,
-            BinaryOperation { .. } | BinaryOperationInplace { .. } | CompareOperation { .. } => -1,
+            BinaryOp { .. } | CompareOperation { .. } => -1,
             BinarySubscript => -1,
             CopyItem { .. } => 1,
             Pop => -1,
@@ -1719,8 +1826,8 @@ impl Instruction {
             DeleteAttr { idx } => w!(DeleteAttr, name = idx),
             LoadConst { idx } => fmt_const("LoadConst", arg, f, idx),
             UnaryOperation { op } => w!(UnaryOperation, ?op),
-            BinaryOperation { op } => w!(BinaryOperation, ?op),
-            BinaryOperationInplace { op } => w!(BinaryOperationInplace, ?op),
+            BinaryOp { op } => write!(f, "{:pad$}({})", "BINARY_OP", op.get(arg)),
+
             BinarySubscript => w!(BinarySubscript),
             LoadAttr { idx } => w!(LoadAttr, name = idx),
             CompareOperation { op } => w!(CompareOperation, ?op),

crates/jit/src/instructions.rs

@@ -436,7 +436,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     _ => Err(JitCompileError::NotSupported),
                 }
             }
-            Instruction::BinaryOperation { op } | Instruction::BinaryOperationInplace { op } => {
+            Instruction::BinaryOp { op } => {
                 let op = op.get(arg);
                 // the rhs is popped off first
                 let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
@@ -446,18 +446,30 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 let b_type = b.to_jit_type();
 
                 let val = match (op, a, b) {
-                    (BinaryOperator::Add, JitValue::Int(a), JitValue::Int(b)) => {
+                    (
+                        BinaryOperator::Add | BinaryOperator::InplaceAdd,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => {
                         let (out, carry) = self.builder.ins().sadd_overflow(a, b);
                         self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW);
                         JitValue::Int(out)
                     }
-                    (BinaryOperator::Subtract, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.compile_sub(a, b))
-                    }
-                    (BinaryOperator::FloorDivide, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().sdiv(a, b))
-                    }
-                    (BinaryOperator::Divide, JitValue::Int(a), JitValue::Int(b)) => {
+                    (
+                        BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.compile_sub(a, b)),
+                    (
+                        BinaryOperator::FloorDivide | BinaryOperator::InplaceFloorDivide,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().sdiv(a, b)),
+                    (
+                        BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => {
                         // Check if b == 0, If so trap with a division by zero error
                         self.builder
                             .ins()
@@ -467,15 +479,21 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                         let b_float = self.builder.ins().fcvt_from_sint(types::F64, b);
                         JitValue::Float(self.builder.ins().fdiv(a_float, b_float))
                     }
-                    (BinaryOperator::Multiply, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().imul(a, b))
-                    }
-                    (BinaryOperator::Modulo, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().srem(a, b))
-                    }
-                    (BinaryOperator::Power, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.compile_ipow(a, b))
-                    }
+                    (
+                        BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().imul(a, b)),
+                    (
+                        BinaryOperator::Remainder | BinaryOperator::InplaceRemainder,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().srem(a, b)),
+                    (
+                        BinaryOperator::Power | BinaryOperator::InplacePower,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.compile_ipow(a, b)),
                     (
                         BinaryOperator::Lshift | BinaryOperator::Rshift,
                         JitValue::Int(a),
@@ -489,39 +507,57 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                             TrapCode::user(CustomTrapCode::NegativeShiftCount as u8).unwrap(),
                         );
 
-                        let out = if op == BinaryOperator::Lshift {
-                            self.builder.ins().ishl(a, b)
-                        } else {
-                            self.builder.ins().sshr(a, b)
-                        };
+                        let out =
+                            if matches!(op, BinaryOperator::Lshift | BinaryOperator::InplaceLshift)
+                            {
+                                self.builder.ins().ishl(a, b)
+                            } else {
+                                self.builder.ins().sshr(a, b)
+                            };
                         JitValue::Int(out)
                     }
-                    (BinaryOperator::And, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().band(a, b))
-                    }
-                    (BinaryOperator::Or, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().bor(a, b))
-                    }
-                    (BinaryOperator::Xor, JitValue::Int(a), JitValue::Int(b)) => {
-                        JitValue::Int(self.builder.ins().bxor(a, b))
-                    }
+                    (
+                        BinaryOperator::And | BinaryOperator::InplaceAnd,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().band(a, b)),
+                    (
+                        BinaryOperator::Or | BinaryOperator::InplaceOr,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().bor(a, b)),
+                    (
+                        BinaryOperator::Xor | BinaryOperator::InplaceXor,
+                        JitValue::Int(a),
+                        JitValue::Int(b),
+                    ) => JitValue::Int(self.builder.ins().bxor(a, b)),
 
                     // Floats
-                    (BinaryOperator::Add, JitValue::Float(a), JitValue::Float(b)) => {
-                        JitValue::Float(self.builder.ins().fadd(a, b))
-                    }
-                    (BinaryOperator::Subtract, JitValue::Float(a), JitValue::Float(b)) => {
-                        JitValue::Float(self.builder.ins().fsub(a, b))
-                    }
-                    (BinaryOperator::Multiply, JitValue::Float(a), JitValue::Float(b)) => {
-                        JitValue::Float(self.builder.ins().fmul(a, b))
-                    }
-                    (BinaryOperator::Divide, JitValue::Float(a), JitValue::Float(b)) => {
-                        JitValue::Float(self.builder.ins().fdiv(a, b))
-                    }
-                    (BinaryOperator::Power, JitValue::Float(a), JitValue::Float(b)) => {
-                        JitValue::Float(self.compile_fpow(a, b))
-                    }
+                    (
+                        BinaryOperator::Add | BinaryOperator::InplaceAdd,
+                        JitValue::Float(a),
+                        JitValue::Float(b),
+                    ) => JitValue::Float(self.builder.ins().fadd(a, b)),
+                    (
+                        BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,
+                        JitValue::Float(a),
+                        JitValue::Float(b),
+                    ) => JitValue::Float(self.builder.ins().fsub(a, b)),
+                    (
+                        BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,
+                        JitValue::Float(a),
+                        JitValue::Float(b),
+                    ) => JitValue::Float(self.builder.ins().fmul(a, b)),
+                    (
+                        BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,
+                        JitValue::Float(a),
+                        JitValue::Float(b),
+                    ) => JitValue::Float(self.builder.ins().fdiv(a, b)),
+                    (
+                        BinaryOperator::Power | BinaryOperator::InplacePower,
+                        JitValue::Float(a),
+                        JitValue::Float(b),
+                    ) => JitValue::Float(self.compile_fpow(a, b)),
 
                     // Floats and Integers
                     (_, JitValue::Int(a), JitValue::Float(b))
@@ -537,19 +573,19 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                         };
 
                         match op {
-                            BinaryOperator::Add => {
+                            BinaryOperator::Add | BinaryOperator::InplaceAdd => {
                                 JitValue::Float(self.builder.ins().fadd(operand_one, operand_two))
                             }
-                            BinaryOperator::Subtract => {
+                            BinaryOperator::Subtract | BinaryOperator::InplaceSubtract => {
                                 JitValue::Float(self.builder.ins().fsub(operand_one, operand_two))
                             }
-                            BinaryOperator::Multiply => {
+                            BinaryOperator::Multiply | BinaryOperator::InplaceMultiply => {
                                 JitValue::Float(self.builder.ins().fmul(operand_one, operand_two))
                             }
-                            BinaryOperator::Divide => {
+                            BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide => {
                                 JitValue::Float(self.builder.ins().fdiv(operand_one, operand_two))
                             }
-                            BinaryOperator::Power => {
+                            BinaryOperator::Power | BinaryOperator::InplacePower => {
                                 JitValue::Float(self.compile_fpow(operand_one, operand_two))
                             }
                             _ => return Err(JitCompileError::NotSupported),

crates/vm/src/frame.rs

@@ -868,10 +868,7 @@ impl ExecutingFrame<'_> {
                 dict.set_item(&*key, value, vm)?;
                 Ok(None)
             }
-            bytecode::Instruction::BinaryOperation { op } => self.execute_bin_op(vm, op.get(arg)),
-            bytecode::Instruction::BinaryOperationInplace { op } => {
-                self.execute_bin_op_inplace(vm, op.get(arg))
-            }
+            bytecode::Instruction::BinaryOp { op } => self.execute_bin_op(vm, op.get(arg)),
             bytecode::Instruction::BinarySubscript => {
                 let key = self.pop_value();
                 let container = self.pop_value();
@@ -2126,40 +2123,29 @@ impl ExecutingFrame<'_> {
             bytecode::BinaryOperator::Multiply => vm._mul(a_ref, b_ref),
             bytecode::BinaryOperator::MatrixMultiply => vm._matmul(a_ref, b_ref),
             bytecode::BinaryOperator::Power => vm._pow(a_ref, b_ref, vm.ctx.none.as_object()),
-            bytecode::BinaryOperator::Divide => vm._truediv(a_ref, b_ref),
+            bytecode::BinaryOperator::TrueDivide => vm._truediv(a_ref, b_ref),
             bytecode::BinaryOperator::FloorDivide => vm._floordiv(a_ref, b_ref),
-            bytecode::BinaryOperator::Modulo => vm._mod(a_ref, b_ref),
+            bytecode::BinaryOperator::Remainder => vm._mod(a_ref, b_ref),
             bytecode::BinaryOperator::Lshift => vm._lshift(a_ref, b_ref),
             bytecode::BinaryOperator::Rshift => vm._rshift(a_ref, b_ref),
             bytecode::BinaryOperator::Xor => vm._xor(a_ref, b_ref),
             bytecode::BinaryOperator::Or => vm._or(a_ref, b_ref),
             bytecode::BinaryOperator::And => vm._and(a_ref, b_ref),
-        }?;
-
-        self.push_value(value);
-        Ok(None)
-    }
-    fn execute_bin_op_inplace(
-        &mut self,
-        vm: &VirtualMachine,
-        op: bytecode::BinaryOperator,
-    ) -> FrameResult {
-        let b_ref = &self.pop_value();
-        let a_ref = &self.pop_value();
-        let value = match op {
-            bytecode::BinaryOperator::Subtract => vm._isub(a_ref, b_ref),
-            bytecode::BinaryOperator::Add => vm._iadd(a_ref, b_ref),
-            bytecode::BinaryOperator::Multiply => vm._imul(a_ref, b_ref),
-            bytecode::BinaryOperator::MatrixMultiply => vm._imatmul(a_ref, b_ref),
-            bytecode::BinaryOperator::Power => vm._ipow(a_ref, b_ref, vm.ctx.none.as_object()),
-            bytecode::BinaryOperator::Divide => vm._itruediv(a_ref, b_ref),
-            bytecode::BinaryOperator::FloorDivide => vm._ifloordiv(a_ref, b_ref),
-            bytecode::BinaryOperator::Modulo => vm._imod(a_ref, b_ref),
-            bytecode::BinaryOperator::Lshift => vm._ilshift(a_ref, b_ref),
-            bytecode::BinaryOperator::Rshift => vm._irshift(a_ref, b_ref),
-            bytecode::BinaryOperator::Xor => vm._ixor(a_ref, b_ref),
-            bytecode::BinaryOperator::Or => vm._ior(a_ref, b_ref),
-            bytecode::BinaryOperator::And => vm._iand(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceSubtract => vm._isub(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceAdd => vm._iadd(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceMultiply => vm._imul(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceMatrixMultiply => vm._imatmul(a_ref, b_ref),
+            bytecode::BinaryOperator::InplacePower => {
+                vm._ipow(a_ref, b_ref, vm.ctx.none.as_object())
+            }
+            bytecode::BinaryOperator::InplaceTrueDivide => vm._itruediv(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceFloorDivide => vm._ifloordiv(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceRemainder => vm._imod(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceLshift => vm._ilshift(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceRshift => vm._irshift(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceXor => vm._ixor(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceOr => vm._ior(a_ref, b_ref),
+            bytecode::BinaryOperator::InplaceAnd => vm._iand(a_ref, b_ref),
         }?;
 
         self.push_value(value);