share comparisonop between bytecode and vm

Author: youknowone

bytecode/src/lib.rs

@@ -264,7 +264,7 @@ pub enum Instruction {
         op: TestOperator,
     },
     CompareOperation {
-        op: CompareOperator,
+        op: ComparisonOperator,
     },
     Pop,
     Rotate2,
@@ -584,17 +584,19 @@ impl<C: Constant> BorrowedConstant<'_, C> {
 }
 
 /// The possible comparison operators
-#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize)]
-pub enum CompareOperator {
-    Greater,
-    GreaterOrEqual,
-    Less,
-    LessOrEqual,
-    Equal,
-    NotEqual,
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ComparisonOperator {
+    // be intentional with bits so that we can do eval_ord with just a bitwise and
+    // bits: | Equal | Greater | Less |
+    Less = 0b001,
+    Greater = 0b010,
+    NotEqual = 0b011,
+    Equal = 0b100,
+    LessOrEqual = 0b101,
+    GreaterOrEqual = 0b110,
 }
 
-#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum TestOperator {
     In,
     NotIn,

compiler/src/compile.rs

@@ -1469,22 +1469,22 @@ impl Compiler {
 
         let compile_cmpop = |op: &ast::Cmpop| match op {
             ast::Cmpop::Eq => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::Equal,
+                op: bytecode::ComparisonOperator::Equal,
             },
             ast::Cmpop::NotEq => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::NotEqual,
+                op: bytecode::ComparisonOperator::NotEqual,
             },
             ast::Cmpop::Lt => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::Less,
+                op: bytecode::ComparisonOperator::Less,
             },
             ast::Cmpop::LtE => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::LessOrEqual,
+                op: bytecode::ComparisonOperator::LessOrEqual,
             },
             ast::Cmpop::Gt => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::Greater,
+                op: bytecode::ComparisonOperator::Greater,
             },
             ast::Cmpop::GtE => Instruction::CompareOperation {
-                op: bytecode::CompareOperator::GreaterOrEqual,
+                op: bytecode::ComparisonOperator::GreaterOrEqual,
             },
             ast::Cmpop::In => Instruction::TestOperation {
                 op: bytecode::TestOperator::In,

jit/src/instructions.rs

@@ -1,8 +1,8 @@
 use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
 use rustpython_bytecode::{
-    self as bytecode, BinaryOperator, BorrowedConstant, CodeObject, CompareOperator, Instruction,
-    Label, UnaryOperator,
+    self as bytecode, BinaryOperator, BorrowedConstant, CodeObject, ComparisonOperator,
+    Instruction, Label, UnaryOperator,
 };
 use std::collections::HashMap;
 
@@ -260,13 +260,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 match (a.ty, b.ty) {
                     (JitType::Int, JitType::Int) => {
                         let cond = match op {
-                            CompareOperator::Equal => IntCC::Equal,
-                            CompareOperator::NotEqual => IntCC::NotEqual,
-                            CompareOperator::Less => IntCC::SignedLessThan,
-                            CompareOperator::LessOrEqual => IntCC::SignedLessThanOrEqual,
-                            CompareOperator::Greater => IntCC::SignedGreaterThan,
-                            CompareOperator::GreaterOrEqual => IntCC::SignedLessThanOrEqual,
-                            _ => return Err(JitCompileError::NotSupported),
+                            ComparisonOperator::Equal => IntCC::Equal,
+                            ComparisonOperator::NotEqual => IntCC::NotEqual,
+                            ComparisonOperator::Less => IntCC::SignedLessThan,
+                            ComparisonOperator::LessOrEqual => IntCC::SignedLessThanOrEqual,
+                            ComparisonOperator::Greater => IntCC::SignedGreaterThan,
+                            ComparisonOperator::GreaterOrEqual => IntCC::SignedLessThanOrEqual,
                         };
 
                         let val = self.builder.ins().icmp(cond, a.val, b.val);

stdlib/src/bisect.rs

@@ -3,7 +3,7 @@ pub(crate) use _bisect::make_module;
 #[pymodule]
 mod _bisect {
     use crate::vm::{
-        function::OptionalArg, types::PyComparisonOp::Lt, PyObjectRef, PyResult, VirtualMachine,
+        function::OptionalArg, types::PyComparisonOp, PyObjectRef, PyResult, VirtualMachine,
     };
 
     #[derive(FromArgs)]
@@ -74,7 +74,9 @@ mod _bisect {
         while lo < hi {
             // Handles issue 13496.
             let mid = (lo + hi) / 2;
-            if a.get_item(mid, vm)?.rich_compare_bool(&x, Lt, vm)? {
+            if a.get_item(mid, vm)?
+                .rich_compare_bool(&x, PyComparisonOp::Lt, vm)?
+            {
                 lo = mid + 1;
             } else {
                 hi = mid;
@@ -102,7 +104,7 @@ mod _bisect {
         while lo < hi {
             // Handles issue 13496.
             let mid = (lo + hi) / 2;
-            if x.rich_compare_bool(&*a.get_item(mid, vm)?, Lt, vm)? {
+            if x.rich_compare_bool(&*a.get_item(mid, vm)?, PyComparisonOp::Lt, vm)? {
                 hi = mid;
             } else {
                 lo = mid + 1;

vm/src/frame.rs

@@ -720,7 +720,7 @@ impl ExecutingFrame<'_> {
             bytecode::Instruction::DeleteAttr { idx } => self.delete_attr(vm, *idx),
             bytecode::Instruction::UnaryOperation { ref op } => self.execute_unop(vm, op),
             bytecode::Instruction::TestOperation { ref op } => self.execute_test(vm, op),
-            bytecode::Instruction::CompareOperation { ref op } => self.execute_compare(vm, op),
+            bytecode::Instruction::CompareOperation { op } => self.execute_compare(vm, *op),
             bytecode::Instruction::ReturnValue => {
                 let value = self.pop_value();
                 self.unwind_blocks(vm, UnwindReason::Returning { value })
@@ -1725,18 +1725,11 @@ impl ExecutingFrame<'_> {
     fn execute_compare(
         &mut self,
         vm: &VirtualMachine,
-        op: &bytecode::CompareOperator,
+        op: bytecode::ComparisonOperator,
     ) -> FrameResult {
         let b = self.pop_value();
         let a = self.pop_value();
-        let op = match *op {
-            bytecode::CompareOperator::Equal => PyComparisonOp::Eq,
-            bytecode::CompareOperator::NotEqual => PyComparisonOp::Ne,
-            bytecode::CompareOperator::Less => PyComparisonOp::Lt,
-            bytecode::CompareOperator::LessOrEqual => PyComparisonOp::Le,
-            bytecode::CompareOperator::Greater => PyComparisonOp::Gt,
-            bytecode::CompareOperator::GreaterOrEqual => PyComparisonOp::Ge,
-        };
+        let op = PyComparisonOp::from(op);
         let value = a.rich_compare(b, op, vm)?;
         self.push_value(value);
         Ok(None)

vm/src/stdlib/operator.rs

@@ -16,47 +16,44 @@ mod _operator {
         function::{ArgBytesLike, FuncArgs, KwArgs, OptionalArg},
         protocol::PyIter,
         recursion::ReprGuard,
-        types::{
-            Callable, Constructor,
-            PyComparisonOp::{Eq, Ge, Gt, Le, Lt, Ne},
-        },
+        types::{Callable, Constructor, PyComparisonOp},
         AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     };
 
     /// Same as a < b.
     #[pyfunction]
     fn lt(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Lt, vm)
+        a.rich_compare(b, PyComparisonOp::Lt, vm)
     }
 
     /// Same as a <= b.
     #[pyfunction]
     fn le(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Le, vm)
+        a.rich_compare(b, PyComparisonOp::Le, vm)
     }
 
     /// Same as a > b.
     #[pyfunction]
     fn gt(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Gt, vm)
+        a.rich_compare(b, PyComparisonOp::Gt, vm)
     }
 
     /// Same as a >= b.
     #[pyfunction]
     fn ge(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Ge, vm)
+        a.rich_compare(b, PyComparisonOp::Ge, vm)
     }
 
     /// Same as a == b.
     #[pyfunction]
     fn eq(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Eq, vm)
+        a.rich_compare(b, PyComparisonOp::Eq, vm)
     }
 
     /// Same as a != b.
     #[pyfunction]
     fn ne(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        a.rich_compare(b, Ne, vm)
+        a.rich_compare(b, PyComparisonOp::Ne, vm)
     }
 
     /// Same as not a.

vm/src/types/slot.rs

@@ -1,6 +1,7 @@
 use crate::common::{hash::PyHash, lock::PyRwLock};
 use crate::{
     builtins::{PyInt, PyStrRef, PyType, PyTypeRef},
+    bytecode::ComparisonOperator,
     convert::{ToPyObject, ToPyResult},
     function::Either,
     function::{FromArgs, FuncArgs, OptionalArg, PyComparisonValue},
@@ -687,18 +688,25 @@ pub trait Comparable: PyPayload {
 }
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
-pub enum PyComparisonOp {
-    // be intentional with bits so that we can do eval_ord with just a bitwise and
-    // bits: | Equal | Greater | Less |
-    Lt = 0b001,
-    Gt = 0b010,
-    Ne = 0b011,
-    Eq = 0b100,
-    Le = 0b101,
-    Ge = 0b110,
-}
-
-use PyComparisonOp::*;
+#[repr(transparent)]
+pub struct PyComparisonOp(ComparisonOperator);
+
+impl From<ComparisonOperator> for PyComparisonOp {
+    fn from(op: ComparisonOperator) -> Self {
+        Self(op)
+    }
+}
+
+#[allow(non_upper_case_globals)]
+impl PyComparisonOp {
+    pub const Lt: Self = Self(ComparisonOperator::Less);
+    pub const Gt: Self = Self(ComparisonOperator::Greater);
+    pub const Ne: Self = Self(ComparisonOperator::NotEqual);
+    pub const Eq: Self = Self(ComparisonOperator::Equal);
+    pub const Le: Self = Self(ComparisonOperator::LessOrEqual);
+    pub const Ge: Self = Self(ComparisonOperator::GreaterOrEqual);
+}
+
 impl PyComparisonOp {
     pub fn eq_only(
         self,
@@ -713,43 +721,43 @@ impl PyComparisonOp {
 
     pub fn eval_ord(self, ord: Ordering) -> bool {
         let bit = match ord {
-            Ordering::Less => Lt,
-            Ordering::Equal => Eq,
-            Ordering::Greater => Gt,
+            Ordering::Less => Self::Lt,
+            Ordering::Equal => Self::Eq,
+            Ordering::Greater => Self::Gt,
         };
-        self as u8 & bit as u8 != 0
+        self.0 as u8 & bit.0 as u8 != 0
     }
 
     pub fn swapped(self) -> Self {
         match self {
-            Lt => Gt,
-            Le => Ge,
-            Eq => Eq,
-            Ne => Ne,
-            Ge => Le,
-            Gt => Lt,
+            Self::Lt => Self::Gt,
+            Self::Le => Self::Ge,
+            Self::Eq => Self::Eq,
+            Self::Ne => Self::Ne,
+            Self::Ge => Self::Le,
+            Self::Gt => Self::Lt,
         }
     }
 
     pub fn method_name(self) -> &'static str {
         match self {
-            Lt => "__lt__",
-            Le => "__le__",
-            Eq => "__eq__",
-            Ne => "__ne__",
-            Ge => "__ge__",
-            Gt => "__gt__",
+            Self::Lt => "__lt__",
+            Self::Le => "__le__",
+            Self::Eq => "__eq__",
+            Self::Ne => "__ne__",
+            Self::Ge => "__ge__",
+            Self::Gt => "__gt__",
         }
     }
 
     pub fn operator_token(self) -> &'static str {
         match self {
-            Lt => "<",
-            Le => "<=",
-            Eq => "==",
-            Ne => "!=",
-            Ge => ">=",
-            Gt => ">",
+            Self::Lt => "<",
+            Self::Le => "<=",
+            Self::Eq => "==",
+            Self::Ne => "!=",
+            Self::Ge => ">=",
+            Self::Gt => ">",
         }
     }