feat: extend stararg fastpath logic to handle lists and generic sequences

Author: BobTheBuidler

mypyc/irbuild/ll_builder.py

@@ -183,7 +183,7 @@
     str_ssize_t_size_op,
     unicode_compare,
 )
-from mypyc.primitives.tuple_ops import list_tuple_op, new_tuple_op, new_tuple_with_length_op
+from mypyc.primitives.tuple_ops import list_tuple_op, new_tuple_op, new_tuple_with_length_op, sequence_tuple_op
 from mypyc.rt_subtype import is_runtime_subtype
 from mypyc.sametype import is_same_type
 from mypyc.subtype import is_subtype
@@ -790,14 +790,22 @@ def _construct_varargs(
                 if star_result is None:
                     # fast path if star expr is a tuple:
                     # we can pass the immutable tuple straight into the function call.
-                    if is_tuple_rprimitive(value.type):
-                        if len(args) == 1:
-                            # fn(*args)
-                            return value, self._create_dict([], [], line)
-                        elif len(args) == 2 and args[1][1] == ARG_STAR2:
-                            # fn(*args, **kwargs)
+                    if len(args) == 1:
+                        # fn(*args)
+                        if is_list_rprimitive(value.type):
+                            value = self.call_c(list_tuple_op, [value], line)
+                        elif not is_tuple_rprimitive(value.type):
+                            value = self.call_c(sequence_tuple_op, [value], line)
+                        return value, self._create_dict([], [], line)
+                    elif len(args) == 2 and args[1][1] == ARG_STAR2:
+                        # fn(*args, **kwargs)
+                        if is_tuple_rprimitive(value.type):
                             star_result = value
-                            continue
+                        elif is_list_rprimitive(value.type):
+                            star_result = self.call_c(list_tuple_op, [value], line)
+                        else:
+                            star_result = self.call_c(sequence_tuple_op, [value], line)
+                        continue
                         # 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)

mypyc/primitives/tuple_ops.py

@@ -76,7 +76,7 @@
 )
 
 # Construct tuple from an arbitrary (iterable) object.
-function_op(
+sequence_tuple_op = function_op(
     name="builtins.tuple",
     arg_types=[object_rprimitive],
     return_type=tuple_rprimitive,

mypyc/test-data/irbuild-basic.test

@@ -3546,3 +3546,267 @@ L0:
     r2 = PyObject_Vectorcall(r1, 0, 0, 0)
     r3 = box(None, 1)
     return r3
+
+[case testStarArgFastPathTuple]
+from typing import Any, Callable
+def deco(fn: Callable[..., Any]) -> Callable[..., Any]:
+    def wrapper(*args: Any) -> Any:
+        return fn(*args)
+    return wrapper
+
+[out]
+def wrapper_deco_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def wrapper_deco_obj.__call__(__mypyc_self__, args):
+    __mypyc_self__ :: __main__.wrapper_deco_obj
+    args :: tuple
+    r0 :: __main__.deco_env
+    r1 :: object
+    r2 :: dict
+    r3 :: object
+L0:
+    r0 = __mypyc_self__.__mypyc_env__
+    r1 = r0.fn
+    r2 = PyDict_New()
+    r3 = PyObject_Call(r1, args, r2)
+    return r3
+def deco(fn):
+    fn :: object
+    r0 :: __main__.deco_env
+    r1 :: bool
+    r2 :: __main__.wrapper_deco_obj
+    r3 :: bool
+    wrapper :: object
+L0:
+    r0 = deco_env()
+    r0.fn = fn; r1 = is_error
+    r2 = wrapper_deco_obj()
+    r2.__mypyc_env__ = r0; r3 = is_error
+    wrapper = r2
+    return wrapper
+
+[case testStarArgFastPathList]
+from typing import Any, Callable, List
+def deco(fn: Callable[..., Any]) -> Callable[[List[Any]], Any]:
+    def wrapper(args: List[Any]) -> Any:
+        return fn(*args)
+    return wrapper
+
+[out]
+def wrapper_deco_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def wrapper_deco_obj.__call__(__mypyc_self__, args):
+    __mypyc_self__ :: __main__.wrapper_deco_obj
+    args :: list
+    r0 :: __main__.deco_env
+    r1 :: object
+    r2 :: tuple
+    r3 :: dict
+    r4 :: object
+L0:
+    r0 = __mypyc_self__.__mypyc_env__
+    r1 = r0.fn
+    r2 = PyList_AsTuple(args)
+    r3 = PyDict_New()
+    r4 = PyObject_Call(r1, r2, r3)
+    return r4
+def deco(fn):
+    fn :: object
+    r0 :: __main__.deco_env
+    r1 :: bool
+    r2 :: __main__.wrapper_deco_obj
+    r3 :: bool
+    wrapper :: object
+L0:
+    r0 = deco_env()
+    r0.fn = fn; r1 = is_error
+    r2 = wrapper_deco_obj()
+    r2.__mypyc_env__ = r0; r3 = is_error
+    wrapper = r2
+    return wrapper
+
+[case testStarArgFastPathListWithKwargs]
+from typing import Any, Callable, Dict, List
+def deco(fn: Callable[..., Any]) -> Callable[..., Any]:
+    def wrapper(lst: List[Any], kwargs: Dict[str, Any]) -> Any:
+        return fn(*lst, **kwargs)
+    return wrapper
+
+[out]
+def wrapper_deco_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def wrapper_deco_obj.__call__(__mypyc_self__, lst, kwargs):
+    __mypyc_self__ :: __main__.wrapper_deco_obj
+    lst :: list
+    kwargs :: dict
+    r0 :: __main__.deco_env
+    r1 :: object
+    r2 :: tuple
+    r3 :: dict
+    r4 :: i32
+    r5 :: bit
+    r6 :: 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
+def deco(fn):
+    fn :: object
+    r0 :: __main__.deco_env
+    r1 :: bool
+    r2 :: __main__.wrapper_deco_obj
+    r3 :: bool
+    wrapper :: object
+L0:
+    r0 = deco_env()
+    r0.fn = fn; r1 = is_error
+    r2 = wrapper_deco_obj()
+    r2.__mypyc_env__ = r0; r3 = is_error
+    wrapper = r2
+    return wrapper
+
+[case testStarArgFastPathSequence]
+from typing import Any, Callable
+def deco(fn: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    def wrapper(args: Any) -> Any:
+        return fn(*args)
+    return wrapper
+
+[out]
+def wrapper_deco_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def wrapper_deco_obj.__call__(__mypyc_self__, args):
+    __mypyc_self__ :: __main__.wrapper_deco_obj
+    args :: object
+    r0 :: __main__.deco_env
+    r1 :: object
+    r2 :: tuple
+    r3 :: dict
+    r4 :: object
+L0:
+    r0 = __mypyc_self__.__mypyc_env__
+    r1 = r0.fn
+    r2 = PySequence_Tuple(args)
+    r3 = PyDict_New()
+    r4 = PyObject_Call(r1, r2, r3)
+    return r4
+def deco(fn):
+    fn :: object
+    r0 :: __main__.deco_env
+    r1 :: bool
+    r2 :: __main__.wrapper_deco_obj
+    r3 :: bool
+    wrapper :: object
+L0:
+    r0 = deco_env()
+    r0.fn = fn; r1 = is_error
+    r2 = wrapper_deco_obj()
+    r2.__mypyc_env__ = r0; r3 = is_error
+    wrapper = r2
+    return wrapper
+
+[case testStarArgFastPathSequenceWithKwargs]
+from typing import Any, Callable
+def deco(fn: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    def wrapper(args: Any, **kwargs: Any) -> Any:
+        return fn(*args, **kwargs)
+    return wrapper
+
+[out]
+def wrapper_deco_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def wrapper_deco_obj.__call__(__mypyc_self__, args, kwargs):
+    __mypyc_self__ :: __main__.wrapper_deco_obj
+    args :: object
+    kwargs :: dict
+    r0 :: __main__.deco_env
+    r1 :: object
+    r2 :: tuple
+    r3 :: dict
+    r4 :: i32
+    r5 :: bit
+    r6 :: 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
+def deco(fn):
+    fn :: object
+    r0 :: __main__.deco_env
+    r1 :: bool
+    r2 :: __main__.wrapper_deco_obj
+    r3 :: bool
+    wrapper :: object
+L0:
+    r0 = deco_env()
+    r0.fn = fn; r1 = is_error
+    r2 = wrapper_deco_obj()
+    r2.__mypyc_env__ = r0; r3 = is_error
+    wrapper = r2
+    return wrapper
+