[mypyc] feat: extend stararg fastpath from #19629 with star2 fastpath

mypyc/irbuild/ll_builder.py

@@ -128,6 +128,8 @@
 from mypyc.primitives.bytes_ops import bytes_compare
 from mypyc.primitives.dict_ops import (
     dict_build_op,
+    dict_copy,
+    dict_copy_op,
     dict_new_op,
     dict_ssize_t_size_op,
     dict_update_in_display_op,
@@ -806,19 +808,40 @@ def _construct_varargs(
                         return value, self._create_dict([], [], line)
                     elif len(args) == 2 and args[1][1] == ARG_STAR2:
                         # fn(*args, **kwargs)
+                        # TODO: extend to cover(*args, **k, **w, **a, **r, **g, **s)
                         if is_tuple_rprimitive(value.type) or isinstance(value.type, RTuple):
                             star_result = value
                         elif is_list_rprimitive(value.type):
                             star_result = self.primitive_op(list_tuple_op, [value], line)
                         else:
                             star_result = self.primitive_op(sequence_tuple_op, [value], line)
-                        continue
+
+                        star2_arg = args[1]
+                        star2_value = star2_arg[0]
+                        if is_dict_rprimitive(star2_value.type):
+                            star2_fastpath_op = dict_copy_op
+                        else:
+                            star2_fastpath_op = dict_copy
+                        return star_result, self.primitive_op(
+                            star2_fastpath_op, [star2_value], line
+                        )
                         # elif ...: TODO extend this to optimize fn(*args, k=1, **kwargs) case
                     # TODO optimize this case using the length utils - currently in review
                     star_result = self.new_list_op(star_values, line)
                 self.primitive_op(list_extend_op, [star_result, value], line)
             elif kind == ARG_STAR2:
                 if star2_result is None:
+                    if len(args) == 1:
+                        # early exit with fastpath if the only arg is ARG_STAR2
+                        # TODO: can we maintain an empty tuple in memory and just reuse it again and again?
+                        if is_dict_rprimitive(value.type):
+                            star2_fastpath_op = dict_copy_op
+                        else:
+                            star2_fastpath_op = dict_copy
+                        return self.new_tuple([], line), self.primitive_op(
+                            star2_fastpath_op, [value], line
+                        )
+
                     star2_result = self._create_dict(star2_keys, star2_values, line)
 
                 self.call_c(dict_update_in_display_op, [star2_result, value], line=line)

mypyc/primitives/dict_ops.py

@@ -53,7 +53,7 @@
 )
 
 # Construct a dictionary from another dictionary.
-function_op(
+dict_copy_op = function_op(
     name="builtins.dict",
     arg_types=[dict_rprimitive],
     return_type=dict_rprimitive,

mypyc/test-data/irbuild-basic.test

@@ -1740,12 +1740,10 @@ def g():
     r6, r7 :: dict
     r8 :: str
     r9 :: object
-    r10 :: dict
-    r11 :: i32
-    r12 :: bit
-    r13 :: tuple
-    r14 :: object
-    r15 :: tuple[int, int, int]
+    r10 :: tuple
+    r11 :: dict
+    r12 :: object
+    r13 :: tuple[int, int, int]
 L0:
     r0 = 'a'
     r1 = 'b'
@@ -1757,13 +1755,11 @@ L0:
     r7 = __main__.globals :: static
     r8 = 'f'
     r9 = CPyDict_GetItem(r7, r8)
-    r10 = PyDict_New()
-    r11 = CPyDict_UpdateInDisplay(r10, r6)
-    r12 = r11 >= 0 :: signed
-    r13 = PyTuple_Pack(0)
-    r14 = PyObject_Call(r9, r13, r10)
-    r15 = unbox(tuple[int, int, int], r14)
-    return r15
+    r10 = PyTuple_Pack(0)
+    r11 = PyDict_Copy(r6)
+    r12 = PyObject_Call(r9, r10, r11)
+    r13 = unbox(tuple[int, int, int], r12)
+    return r13
 def h():
     r0, r1 :: str
     r2, r3 :: object
@@ -3670,18 +3666,14 @@ def wrapper_deco_obj.__call__(__mypyc_self__, lst, kwargs):
     r1 :: object
     r2 :: tuple
     r3 :: dict
-    r4 :: i32
-    r5 :: bit
-    r6 :: object
+    r4 :: object
 L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PyList_AsTuple(lst)
-    r3 = PyDict_New()
-    r4 = CPyDict_UpdateInDisplay(r3, kwargs)
-    r5 = r4 >= 0 :: signed
-    r6 = PyObject_Call(r1, r2, r3)
-    return r6
+    r3 = PyDict_Copy(kwargs)
+    r4 = PyObject_Call(r1, r2, r3)
+    return r4
 def deco(fn):
     fn :: object
     r0 :: __main__.deco_env
@@ -3777,18 +3769,14 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args, kwargs):
     r1 :: object
     r2 :: tuple
     r3 :: dict
-    r4 :: i32
-    r5 :: bit
-    r6 :: object
+    r4 :: object
 L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = r0.fn
     r2 = PySequence_Tuple(args)
-    r3 = PyDict_New()
-    r4 = CPyDict_UpdateInDisplay(r3, kwargs)
-    r5 = r4 >= 0 :: signed
-    r6 = PyObject_Call(r1, r2, r3)
-    return r6
+    r3 = PyDict_Copy(kwargs)
+    r4 = PyObject_Call(r1, r2, r3)
+    return r4
 def deco(fn):
     fn :: object
     r0 :: __main__.deco_env

mypyc/test-data/irbuild-generics.test

@@ -167,17 +167,13 @@ def execute(func, args, kwargs):
     func :: object
     args :: tuple
     kwargs, r0 :: dict
-    r1 :: i32
-    r2 :: bit
-    r3 :: object
-    r4 :: int
+    r1 :: object
+    r2 :: int
 L0:
-    r0 = PyDict_New()
-    r1 = CPyDict_UpdateInDisplay(r0, kwargs)
-    r2 = r1 >= 0 :: signed
-    r3 = PyObject_Call(func, args, r0)
-    r4 = unbox(int, r3)
-    return r4
+    r0 = PyDict_Copy(kwargs)
+    r1 = PyObject_Call(func, args, r0)
+    r2 = unbox(int, r1)
+    return r2
 def f(x):
     x :: int
 L0:
@@ -703,10 +699,8 @@ def inner_deco_obj.__call__(__mypyc_self__, args, kwargs):
     r22, can_iter, r23, can_use_keys, r24, can_use_values :: list
     r25 :: object
     r26 :: dict
-    r27 :: i32
-    r28 :: bit
-    r29 :: object
-    r30 :: int
+    r27 :: object
+    r28 :: int
 L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = var_object_size args
@@ -758,12 +752,10 @@ L9:
     r24 = CPyDict_Values(kwargs)
     can_use_values = r24
     r25 = r0.func
-    r26 = PyDict_New()
-    r27 = CPyDict_UpdateInDisplay(r26, kwargs)
-    r28 = r27 >= 0 :: signed
-    r29 = PyObject_Call(r25, args, r26)
-    r30 = unbox(int, r29)
-    return r30
+    r26 = PyDict_Copy(kwargs)
+    r27 = PyObject_Call(r25, args, r26)
+    r28 = unbox(int, r27)
+    return r28
 def deco(func):
     func :: object
     r0 :: __main__.deco_env

mypyc/test-data/run-functions.test

@@ -1312,3 +1312,29 @@ from native import f
 print(f(1))
 [out]
 2
+
+[case testStarArgFastPaths]
+from typing import Any, Mapping
+def fn(x: str, y: int) -> str:
+    return x * y
+def star_tuple(*args: Any) -> str:
+    return fn(*args)
+def star_list(args: list[Any]) -> str:
+    return fn(*args)
+def star_generic(args: dict[Any, Any]) -> str:
+    return fn(*args)
+def star2(**kwargs: Any) -> str:
+    return fn(**kwargs)
+def star2_generic(kwargs: Mapping[Any, Any]) -> str:
+    return fn(**kwargs)
+
+def test_star_fastpath_tuple() -> None:
+    assert star_tuple("a", 3) == "aaa"
+def test_star_fastpath_list() -> None:
+    assert star_list(["a", 3]) == "aaa"
+def test_star_fastpath_generic() -> None:
+    assert star_generic({"a": None, 3: None}) == "aaa"
+def test_star2_fastpath() -> None:
+    assert star2(x="a", y=3) == "aaa"
+def test_star2_fastpath_generic() -> None:
+    assert star2_generic({"x": "a", "y": 3}) == "aaa"