[mypyc] Speed up "in" against tuple

Also make the semantics closer to Python.

Author: JukkaL

mypyc/irbuild/expression.py

@@ -752,44 +752,60 @@ def go(i: int, prev: Value) -> Value:
 def try_specialize_in_expr(
     builder: IRBuilder, op: str, lhs: Expression, rhs: Expression, line: int
 ) -> Value | None:
-    if isinstance(rhs, (TupleExpr, ListExpr)):
-        items = rhs.items
+    left: Value | None = None
+    items: list[Value] | None = None
+
+    # 16 is arbitrarily chosen to limit code size
+    if isinstance(rhs, (TupleExpr, ListExpr)) and len(rhs.items) < 16:
+        left = builder.accept(lhs)
+        items = [builder.accept(item) for item in rhs.items]
+    elif isinstance(builder.node_type(rhs), RTuple):
+        left = builder.accept(lhs)
+        tuple_val = builder.accept(rhs)
+        assert isinstance(tuple_val.type, RTuple)
+        items = [builder.add(TupleGet(tuple_val, i)) for i in range(len(tuple_val.type.types))]
+
+    if items is not None:
+        assert left is not None
         n_items = len(items)
         # x in y -> x == y[0] or ... or x == y[n]
         # x not in y -> x != y[0] and ... and x != y[n]
-        # 16 is arbitrarily chosen to limit code size
-        if 1 < n_items < 16:
+        if n_items > 1:
             if op == "in":
-                bin_op = "or"
                 cmp_op = "=="
             else:
-                bin_op = "and"
                 cmp_op = "!="
-            mypy_file = builder.graph["builtins"].tree
-            assert mypy_file is not None
-            info = mypy_file.names["bool"].node
-            assert isinstance(info, TypeInfo), info
-            bool_type = Instance(info, [])
-            exprs = []
+            out = BasicBlock()
             for item in items:
-                expr = ComparisonExpr([cmp_op], [lhs, item])
-                builder.types[expr] = bool_type
-                exprs.append(expr)
-
-            or_expr: Expression = exprs.pop(0)
-            for expr in exprs:
-                or_expr = OpExpr(bin_op, or_expr, expr)
-                builder.types[or_expr] = bool_type
-            return builder.accept(or_expr)
+                x = transform_basic_comparison(builder, cmp_op, left, item, line)
+                b = builder.builder.bool_value(x)
+                nxt = BasicBlock()
+                if op == "in":
+                    builder.add_bool_branch(b, out, nxt)
+                else:
+                    builder.add_bool_branch(b, nxt, out)
+                builder.activate_block(nxt)
+            r = Register(bool_rprimitive)
+            end = BasicBlock()
+            if op == "in":
+                values = builder.false(), builder.true()
+            else:
+                values = builder.true(), builder.false()
+            builder.assign(r, values[0], line)
+            builder.goto(end)
+            builder.activate_block(out)
+            builder.assign(r, values[1], line)
+            builder.goto(end)
+            builder.activate_block(end)
+            return r
         # x in [y]/(y) -> x == y
         # x not in [y]/(y) -> x != y
         elif n_items == 1:
             if op == "in":
                 cmp_op = "=="
             else:
                 cmp_op = "!="
-            left = builder.accept(lhs)
-            right = builder.accept(items[0])
+            right = items[0]
             return transform_basic_comparison(builder, cmp_op, left, right, line)
         # x in []/() -> False
         # x not in []/() -> True

mypyc/test-data/irbuild-tuple.test

@@ -184,30 +184,102 @@ def f(i: int) -> bool:
 [out]
 def f(i):
     i :: int
-    r0 :: bit
-    r1 :: bool
-    r2 :: bit
+    r0, r1, r2 :: bit
     r3 :: bool
-    r4 :: bit
 L0:
     r0 = int_eq i, 2
-    if r0 goto L1 else goto L2 :: bool
+    if r0 goto L4 else goto L1 :: bool
 L1:
-    r1 = r0
-    goto L3
+    r1 = int_eq i, 4
+    if r1 goto L4 else goto L2 :: bool
+L2:
+    r2 = int_eq i, 6
+    if r2 goto L4 else goto L3 :: bool
+L3:
+    r3 = 0
+    goto L5
+L4:
+    r3 = 1
+L5:
+    return r3
+
+[case testTupleOperatorNotIn]
+def x() -> int:
+    return 1
+def y() -> int:
+    return 2
+def z() -> int:
+    return 3
+
+def f() -> bool:
+    return z() not in (x(), y())
+[out]
+def x():
+L0:
+    return 2
+def y():
+L0:
+    return 4
+def z():
+L0:
+    return 6
+def f():
+    r0, r1, r2 :: int
+    r3, r4 :: bit
+    r5 :: bool
+L0:
+    r0 = z()
+    r1 = x()
+    r2 = y()
+    r3 = int_ne r0, r1
+    if r3 goto L1 else goto L3 :: bool
+L1:
+    r4 = int_ne r0, r2
+    if r4 goto L2 else goto L3 :: bool
+L2:
+    r5 = 1
+    goto L4
+L3:
+    r5 = 0
+L4:
+    return r5
+
+[case testTupleOperatorInFinalTuple]
+from typing import Final
+
+tt: Final = (1, 2)
+
+def f(x: int) -> bool:
+    return x in tt
+[out]
+def f(x):
+    x :: int
+    r0 :: tuple[int, int]
+    r1 :: bool
+    r2, r3 :: int
+    r4, r5 :: bit
+    r6 :: bool
+L0:
+    r0 = __main__.tt :: static
+    if is_error(r0) goto L1 else goto L2
+L1:
+    r1 = raise NameError('value for final name "tt" was not set')
+    unreachable
 L2:
-    r2 = int_eq i, 4
-    r1 = r2
+    r2 = r0[0]
+    r3 = r0[1]
+    r4 = int_eq x, r2
+    if r4 goto L5 else goto L3 :: bool
 L3:
-    if r1 goto L4 else goto L5 :: bool
+    r5 = int_eq x, r3
+    if r5 goto L5 else goto L4 :: bool
 L4:
-    r3 = r1
+    r6 = 0
     goto L6
 L5:
-    r4 = int_eq i, 6
-    r3 = r4
+    r6 = 1
 L6:
-    return r3
+    return r6
 
 [case testTupleOperatorInFinalTuple]
 from typing import Final