Fix: infer_BINARY_OP now supports smallint input (python#46)

JuliaPoo · Fidget-Spinner · web-flow · commit 46f6db41afad · 2023-06-25T23:47:19.000+08:00
Co-authored-by: Ken Jin &lt;kenjin4096@gmail.com&gt;
diff --git a/Python/tier2.c b/Python/tier2.c
@@ -1616,12 +1616,13 @@ infer_BINARY_OP(
     PyTypeObject *righttype = (PyTypeObject *)_Py_TYPENODE_CLEAR_TAG(rightroot);
 
     if (_Py_TYPENODE_IS_POSITIVE_NULL(leftroot)
-        && (righttype == &PyLong_Type
+        && (righttype == &PyLong_Type || righttype == &PySmallInt_Type
             || righttype == &PyFloat_Type || righttype == &PyRawFloat_Type)) {
         // Check if same type as right
         *needs_guard = true;
         write_curr = emit_type_guard(write_curr,
-            righttype == &PyLong_Type ? CHECK_INT : CHECK_FLOAT, 1, bb_id);
+            (righttype == &PyLong_Type || righttype == &PySmallInt_Type)
+            ? CHECK_INT : CHECK_FLOAT, 1, bb_id);
         return write_curr;
     }
     if (_Py_TYPENODE_GET_TAG(leftroot) == TYPE_ROOT_NEGATIVE) {
@@ -1661,7 +1662,8 @@ infer_BINARY_OP(
         type_propagate(opcode, 0, type_context, NULL);
         return write_curr;
     }
-    if (righttype == &PyLong_Type && lefttype == &PyLong_Type) {
+    if ((righttype == &PyLong_Type || righttype == &PySmallInt_Type)
+        && (lefttype == &PyLong_Type || lefttype == &PySmallInt_Type)) {
         int opcode = oparg == NB_ADD
             ? BINARY_OP_ADD_INT_REST
             : oparg == NB_SUBTRACT
diff --git a/README.md b/README.md
@@ -138,6 +138,7 @@ We provide pre-compiled binaries for 64-bit Windows 10/11 via GitHub releases.
     * Mark interpreter frames as tier 2 or not in https://github.com/pylbbv/pylbbv/pull/34
     * Make instruction offset calculation frame aware in https://github.com/pylbbv/pylbbv/pull/37
     * Support multiple entry points in a Basic Block in https://github.com/pylbbv/pylbbv/pull/39 and https://github.com/pylbbv/pylbbv/pull/40
+    * `BINARY_OP` specialisations now apply to `PySmallInt_Type` in https://github.com/pylbbv/pylbbv/pull/46
 * Improved workflow:
     * workflow: enable CI GH actions tests in https://github.com/pylbbv/pylbbv/pull/35
 * Improved type propagator:
diff --git a/tier2_test.py b/tier2_test.py
@@ -615,6 +615,36 @@ def f(x, items):
 
     # As long as it doesn't crash, everything's good
 
+with TestInfo("infer_BINARY_OP to be compatible with smallint"):
+    # See https://github.com/pylbbv/pylbbv/issues/45 for more information.
+
+    # Testing left smallint
+    def f(a, b):
+        z = a + b
+        return z + 1
+
+    trigger_tier2(f, (1,1))
+    insts = dis.get_instructions(test_typeprop1, tier2=True)
+    assert [x.opname for x in insts].count("BINARY_OP_ADD_INT_REST") == 2
+
+    # Testing right smallint
+    def f(a, b):
+        z = a + b
+        return 1 + z
+
+    trigger_tier2(f, (1,1))
+    insts = dis.get_instructions(test_typeprop1, tier2=True)
+    assert [x.opname for x in insts].count("BINARY_OP_ADD_INT_REST") == 2
+
+    # Testing both sides smallint
+    def f(a, b):
+        z = a + b
+        w = 1
+        return 1 + w
+    trigger_tier2(f, (1,1))
+    insts = dis.get_instructions(test_typeprop1, tier2=True)
+    assert [x.opname for x in insts].count("BINARY_OP_ADD_INT_REST") == 2
+
 print("Regression tests...Done!")
 
 print("Tests completed ^-^")
