Repository: cython/cython
Pull request: Optimize numerical operations with python integer variable
For now let's get optimisation of int variable on par with untype variable. Further optimisations can be done in separate PR.

Fixes #5783.


Problem statement:
[BUG] Operations with variable annotated with `int` avoids optimisations
### Describe the bug

When variable has type `int` in pure python mode, Cython is using generic `PyObject_RichCompare` instead of optimised `__Pyx_PyInt_EqObjC` which is used when no annotation is used.

### Code to reproduce the behaviour:

```python
def bar(a: int):
    print(a == 5)

bar()
```

Generated C code:

```c
  /* "deleteme.py":23
 * 
 * def bar(a: int):
 *     print(a == 5)             # <<<<<<<<<<<<<<
 * 
 * bar()
 */
  __pyx_t_1 = PyObject_RichCompare(__pyx_v_a, __pyx_int_5, Py_EQ); __Pyx_XGOTREF(__pyx_t_1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 23, __pyx_L1_error)
```

### Expected behaviour

```python
def bar(a):
    print(a == 5)

bar()
```

Generated C code:
```c
  /* "deleteme.py":23
 * 
 * def bar(a):
 *     print(a == 5)             # <<<<<<<<<<<<<<
 * 
 * bar()
 */
  __pyx_t_1 = __Pyx_PyInt_EqObjC(__pyx_v_a, __pyx_int_5, 5, 0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 23, __pyx_L1_error)
  __Pyx_GOTREF(__pyx_t_1);
```

### OS

MacOS

### Python version

3.9

### Cython version

master

### Additional context

_No response_

Hints from issue discussion:
I suspect what's happened is that the check is probably based on `is py_object_type` and we assume anything typed with something else *isn't* an int.

We could probably do better though if we really wanted:
`__Pyx_PyInt_EqObjC` is written with the assumption that `a` is probably an `int` but doesn't have to be. So we could generate a more specific function for this.
> I suspect what's happened is that the check is probably based on is py_object_type and we assume anything typed with something else isn't an int.

Yes there are two places:
1. New method needs to be added for `int` type - `_handle_simple_method_int___eq__`:
 https://github.com/cython/cython/blob/edaea23cc5a7909e84d3bd01b0a8ab2206af8e95/Cython/Compiler/Optimize.py#L3425
2. `int` type needs to be added to this line:
 https://github.com/cython/cython/blob/edaea23cc5a7909e84d3bd01b0a8ab2206af8e95/Cython/Compiler/Optimize.py#L4210
https://github.com/cython/cython/blob/edaea23cc5a7909e84d3bd01b0a8ab2206af8e95/Cython/Compiler/Optimize.py#L4215

> 1. New method needs to be added for `int` type - `_handle_simple_method_int___eq__`:

I'm sure you've spotted it, but this issue (and fix) applies to most other binops too (i.e. add, sub, etc)

Commit message: Assert tree in tests

Patch to review:
diff --git a/Cython/Compiler/Optimize.py b/Cython/Compiler/Optimize.py
index 42fab2e2cdb..5138e016a3c 100644
--- a/Cython/Compiler/Optimize.py
+++ b/Cython/Compiler/Optimize.py
@@ -3463,6 +3463,20 @@ def _handle_simple_method_object___truediv__(self, node, function, args, is_unbo
     def _handle_simple_method_object___div__(self, node, function, args, is_unbound_method):
         return self._optimise_num_div('Divide', node, function, args, is_unbound_method)
 
+    _handle_simple_method_int___add__ = _handle_simple_method_object___add__
+    _handle_simple_method_int___sub__ = _handle_simple_method_object___sub__
+    _handle_simple_method_int___mul__ = _handle_simple_method_object___mul__
+    _handle_simple_method_int___eq__ = _handle_simple_method_object___eq__
+    _handle_simple_method_int___ne__ = _handle_simple_method_object___ne__
+    _handle_simple_method_int___and__ = _handle_simple_method_object___and__
+    _handle_simple_method_int___or__ = _handle_simple_method_object___or__
+    _handle_simple_method_int___xor__ = _handle_simple_method_object___xor__
+    _handle_simple_method_int___rshift__ = _handle_simple_method_object___rshift__
+    _handle_simple_method_int___lshift__ = _handle_simple_method_object___lshift__
+    _handle_simple_method_int___mod__ = _handle_simple_method_object___mod__
+    _handle_simple_method_int___floordiv__ = _handle_simple_method_object___floordiv__
+    _handle_simple_method_int___truediv__ = _handle_simple_method_object___truediv__
+
     def _optimise_num_div(self, operator, node, function, args, is_unbound_method):
         if len(args) != 2 or not args[1].has_constant_result() or args[1].constant_result == 0:
             return node
@@ -4207,12 +4221,12 @@ def optimise_numeric_binop(operator, node, ret_type, arg0, arg1):
     # Prefer constants on RHS as they allows better size control for some operators.
     num_nodes = (ExprNodes.IntNode, ExprNodes.FloatNode)
     if isinstance(arg1, num_nodes):
-        if arg0.type is not PyrexTypes.py_object_type:
+        if arg0.type is not PyrexTypes.py_object_type and arg0.type is not Builtin.builtin_types['int']:
             return None
         numval = arg1
         arg_order = 'ObjC'
     elif isinstance(arg0, num_nodes):
-        if arg1.type is not PyrexTypes.py_object_type:
+        if arg1.type is not PyrexTypes.py_object_type and arg1.type is not Builtin.builtin_types['int']:
             return None
         numval = arg0
         arg_order = 'CObj'
diff --git a/tests/run/pyintop.pyx b/tests/run/pyintop.pyx
index e1b93829212..0ba3c86a65a 100644
--- a/tests/run/pyintop.pyx
+++ b/tests/run/pyintop.pyx
@@ -1,3 +1,4 @@
+# cython: language_level=3
 # mode: run
 
 cimport cython
@@ -515,3 +516,158 @@ def test_avoid_if_coercion(obj):
         return True
     else:
         return False
+
+@cython.test_fail_if_path_exists('//AddNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_add_int(obj2: int):
+    """
+    >>> pure_add_int(1)
+    (2, 2)
+    """
+    res1 = obj2 + 1
+    res2 = 1 + obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//SubNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_sub_int(obj2: int):
+    """
+    >>> pure_sub_int(1)
+    (0, 0)
+    """
+    res1 = obj2 - 1
+    res2 = 1 - obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//MulNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_mul_int(obj2: int):
+    """
+    >>> pure_mul_int(2)
+    (4, 4)
+    """
+    res1 = obj2 * 2
+    res2 = 2 * obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//PrimaryCmpNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_eq_int(obj2: int):
+    """
+    >>> pure_eq_int(2)
+    (True, True)
+    >>> pure_eq_int(3)
+    (False, False)
+    """
+    res1 = obj2 == 2
+    res2 = 2 == obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//PrimaryCmpNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_ne_int(obj2: int):
+    """
+    >>> pure_ne_int(2)
+    (False, False)
+    >>> pure_ne_int(3)
+    (True, True)
+    """
+    res1 = obj2 != 2
+    res2 = 2 != obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_and_int(obj2: int):
+    """
+    >>> pure_and_int(1)
+    (0, 0)
+    >>> pure_and_int(3)
+    (2, 2)
+    """
+    res1 = obj2 & 2
+    res2 = 2 & obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_or_int(obj2: int):
+    """
+    >>> pure_or_int(1)
+    (3, 3)
+    >>> pure_or_int(0)
+    (2, 2)
+    """
+    res1 = obj2 | 2
+    res2 = 2 | obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_xor_int(obj2: int):
+    """
+    >>> pure_xor_int(1)
+    (3, 3)
+    >>> pure_xor_int(3)
+    (1, 1)
+    """
+    res1 = obj2 ^ 2
+    res2 = 2 ^ obj2
+    return res1, res2
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_rshift_int(obj2: int):
+    """
+    >>> pure_rshift_int(8)
+    4
+    """
+    res = obj2 >> 1
+    return res
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_lshift_int(obj2: int):
+    """
+    >>> pure_lshift_int(8)
+    16
+    """
+    res = obj2 << 1
+    return res
+
+@cython.test_fail_if_path_exists('//IntBinopNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_mod_int(obj2: int):
+    """
+    >>> pure_mod_int(3)
+    1
+    """
+    res = obj2 % 2
+    return res
+
+@cython.test_fail_if_path_exists('//DivNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_floordiv_int(obj2: int):
+    """
+    >>> pure_floordiv_int(3)
+    1
+    """
+    res = obj2 // 2
+    return res
+
+import sys
+
+
+@cython.test_fail_if_path_exists('//DivNode')
+@cython.test_assert_path_exists('//PythonCapiFunctionNode')
+def pure_truediv_int(obj2: int):
+    """
+    >>> pure_truediv_int(5)
+    2.5
+    """
+    res = obj2 / 2
+    if sys.version_info >= (3, 0):
+        return res
+    else:
+        # Workaround for python 2.7 to skip this test.
+        return 2.5


Code review comments:
diff_hunk: @@ -4207,12 +4221,12 @@ def optimise_numeric_binop(operator, node, ret_type, arg0, arg1):
     # Prefer constants on RHS as they allows better size control for some operators.
     num_nodes = (ExprNodes.IntNode, ExprNodes.FloatNode)
     if isinstance(arg1, num_nodes):
-        if arg0.type is not PyrexTypes.py_object_type:
+        if arg0.type is not PyrexTypes.py_object_type and arg0.type is not Builtin.builtin_types['int']:
             return None
         numval = arg1
         arg_order = 'ObjC'
     elif isinstance(arg0, num_nodes):
-        if arg1.type is not PyrexTypes.py_object_type:
+        if arg1.type is not PyrexTypes.py_object_type and arg1.type is not Builtin.builtin_types['int']:
path: Cython/Compiler/Optimize.py
text: @user1:
```suggestion
        if arg1.type is not PyrexTypes.py_object_type and arg1.type is not Builtin.int_type:
```

diff_hunk: @@ -4207,12 +4221,12 @@ def optimise_numeric_binop(operator, node, ret_type, arg0, arg1):
     # Prefer constants on RHS as they allows better size control for some operators.
     num_nodes = (ExprNodes.IntNode, ExprNodes.FloatNode)
     if isinstance(arg1, num_nodes):
-        if arg0.type is not PyrexTypes.py_object_type:
+        if arg0.type is not PyrexTypes.py_object_type and arg0.type is not Builtin.builtin_types['int']:
path: Cython/Compiler/Optimize.py
text: @user1:
Not sure if we need the Python `long` type here as well. That might not settle itself before Cython 3.1…
Annotations will probably use `int`, so let's stick to that.
```suggestion
        if arg0.type is not PyrexTypes.py_object_type and arg0.type is not Builtin.int_type:
```

diff_hunk: @@ -1,3 +1,4 @@
+# cython: language_level=3
path: tests/run/pyintop.pyx
text: @user1:
Let's just use the new default. No need to nail the string  type as well.
```suggestion
# cython: language_level=3str
```