[mypyc] Add primitive ops for dictionary methods (#8742)

This covers both view and list versions (it looks like both are used relatively often, at least in mypy). This accompanies #8725 to account for cases where `keys`/`values`/`items` appear in non-loop contexts.

Author: ilevkivskyi

mypyc/irbuild/specialize.py

@@ -22,8 +22,9 @@
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive
+    bool_rprimitive, is_dict_rprimitive
 )
+from mypyc.primitives.dict_ops import dict_keys_op, dict_values_op, dict_items_op
 from mypyc.primitives.misc_ops import true_op, false_op
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.for_helpers import translate_list_comprehension, comprehension_helper
@@ -77,6 +78,34 @@ def translate_len(
     return None
 
 
+@specialize_function('builtins.list')
+def dict_methods_fast_path(
+        builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
+    # Specialize a common case when list() is called on a dictionary view
+    # method call, for example foo = list(bar.keys()).
+    if not (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]):
+        return None
+    arg = expr.args[0]
+    if not (isinstance(arg, CallExpr) and not arg.args
+            and isinstance(arg.callee, MemberExpr)):
+        return None
+    base = arg.callee.expr
+    attr = arg.callee.name
+    rtype = builder.node_type(base)
+    if not (is_dict_rprimitive(rtype) and attr in ('keys', 'values', 'items')):
+        return None
+
+    obj = builder.accept(base)
+    # Note that it is not safe to use fast methods on dict subclasses, so
+    # the corresponding helpers in CPy.h fallback to (inlined) generic logic.
+    if attr == 'keys':
+        return builder.primitive_op(dict_keys_op, [obj], expr.line)
+    elif attr == 'values':
+        return builder.primitive_op(dict_values_op, [obj], expr.line)
+    else:
+        return builder.primitive_op(dict_items_op, [obj], expr.line)
+
+
 @specialize_function('builtins.tuple')
 @specialize_function('builtins.set')
 @specialize_function('builtins.dict')

mypyc/lib-rt/CPy.h

@@ -1377,6 +1377,84 @@ static tuple_T3OOO CPy_GetExcInfo(void) {
     return ret;
 }
 
+static PyObject *CPyDict_KeysView(PyObject *dict) {
+    if (PyDict_CheckExact(dict)){
+        return _CPyDictView_New(dict, &PyDictKeys_Type);
+    }
+    return PyObject_CallMethod(dict, "keys", NULL);
+}
+
+static PyObject *CPyDict_ValuesView(PyObject *dict) {
+    if (PyDict_CheckExact(dict)){
+        return _CPyDictView_New(dict, &PyDictValues_Type);
+    }
+    return PyObject_CallMethod(dict, "values", NULL);
+}
+
+static PyObject *CPyDict_ItemsView(PyObject *dict) {
+    if (PyDict_CheckExact(dict)){
+        return _CPyDictView_New(dict, &PyDictItems_Type);
+    }
+    return PyObject_CallMethod(dict, "items", NULL);
+}
+
+static PyObject *CPyDict_Keys(PyObject *dict) {
+    if PyDict_CheckExact(dict) {
+        return PyDict_Keys(dict);
+    }
+    // Inline generic fallback logic to also return a list.
+    PyObject *list = PyList_New(0);
+    PyObject *view = PyObject_CallMethod(dict, "keys", NULL);
+    if (view == NULL) {
+        return NULL;
+    }
+    PyObject *res = _PyList_Extend((PyListObject *)list, view);
+    Py_DECREF(view);
+    if (res == NULL) {
+        return NULL;
+    }
+    Py_DECREF(res);
+    return list;
+}
+
+static PyObject *CPyDict_Values(PyObject *dict) {
+    if PyDict_CheckExact(dict) {
+        return PyDict_Values(dict);
+    }
+    // Inline generic fallback logic to also return a list.
+    PyObject *list = PyList_New(0);
+    PyObject *view = PyObject_CallMethod(dict, "values", NULL);
+    if (view == NULL) {
+        return NULL;
+    }
+    PyObject *res = _PyList_Extend((PyListObject *)list, view);
+    Py_DECREF(view);
+    if (res == NULL) {
+        return NULL;
+    }
+    Py_DECREF(res);
+    return list;
+}
+
+static PyObject *CPyDict_Items(PyObject *dict) {
+    if PyDict_CheckExact(dict) {
+        return PyDict_Items(dict);
+    }
+    // Inline generic fallback logic to also return a list.
+    PyObject *list = PyList_New(0);
+    PyObject *view = PyObject_CallMethod(dict, "items", NULL);
+    if (view == NULL) {
+        return NULL;
+    }
+    PyObject *res = _PyList_Extend((PyListObject *)list, view);
+    Py_DECREF(view);
+    if (res == NULL) {
+        return NULL;
+    }
+    Py_DECREF(res);
+    return list;
+}
+
 static PyObject *CPyDict_GetKeysIter(PyObject *dict) {
     if (PyDict_CheckExact(dict)) {
         // Return dict itself to indicate we can use fast path instead.

mypyc/lib-rt/pythonsupport.h

@@ -366,6 +366,25 @@ CPyGen_SetStopIterationValue(PyObject *value)
     return 0;
 }
 
+// Copied from dictobject.c and dictobject.h, these are not Public before
+// Python 3.8. Also remove some error checks that we do in the callers.
+typedef struct {
+    PyObject_HEAD
+    PyDictObject *dv_dict;
+} _CPyDictViewObject;
+
+static PyObject *
+_CPyDictView_New(PyObject *dict, PyTypeObject *type)
+{
+    _CPyDictViewObject *dv = PyObject_GC_New(_CPyDictViewObject, type);
+    if (dv == NULL)
+        return NULL;
+    Py_INCREF(dict);
+    dv->dv_dict = (PyDictObject *)dict;
+    PyObject_GC_Track(dv);
+    return (PyObject *)dv;
+}
+
 #ifdef __cplusplus
 }
 #endif

mypyc/primitives/dict_ops.py

@@ -5,7 +5,7 @@
 from mypyc.ir.ops import EmitterInterface, ERR_FALSE, ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
     dict_rprimitive, object_rprimitive, bool_rprimitive, int_rprimitive,
-    dict_next_rtuple_single, dict_next_rtuple_pair
+    list_rprimitive, dict_next_rtuple_single, dict_next_rtuple_pair
 )
 
 from mypyc.primitives.registry import (
@@ -125,6 +125,60 @@ def emit_new_dict(emitter: EmitterInterface, args: List[str], dest: str) -> None
     error_kind=ERR_MAGIC,
     emit=call_emit('CPyDict_FromAny'))
 
+# dict.keys()
+method_op(
+    name='keys',
+    arg_types=[dict_rprimitive],
+    result_type=object_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_KeysView')
+)
+
+# dict.values()
+method_op(
+    name='values',
+    arg_types=[dict_rprimitive],
+    result_type=object_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_ValuesView')
+)
+
+# dict.items()
+method_op(
+    name='items',
+    arg_types=[dict_rprimitive],
+    result_type=object_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_ItemsView')
+)
+
+# list(dict.keys())
+dict_keys_op = custom_op(
+    name='keys',
+    arg_types=[dict_rprimitive],
+    result_type=list_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_Keys')
+)
+
+# list(dict.values())
+dict_values_op = custom_op(
+    name='values',
+    arg_types=[dict_rprimitive],
+    result_type=list_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_Values')
+)
+
+# list(dict.items())
+dict_items_op = custom_op(
+    name='items',
+    arg_types=[dict_rprimitive],
+    result_type=list_rprimitive,
+    error_kind=ERR_MAGIC,
+    emit=call_emit('CPyDict_Items')
+)
+
 
 def emit_len(emitter: EmitterInterface, args: List[str], dest: str) -> None:
     temp = emitter.temp_name()

mypyc/test-data/fixtures/typing-full.pyi

@@ -131,6 +131,7 @@ class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta):
     @overload
     def get(self, k: T, default: Union[T_co, V]) -> Union[T_co, V]: pass
     def values(self) -> Iterable[T_co]: pass  # Approximate return type
+    def items(self) -> Iterable[Tuple[T, T_co]]: pass  # Approximate return type
     def __len__(self) -> int: ...
     def __contains__(self, arg: object) -> int: pass
 

mypyc/test-data/run.test

@@ -989,7 +989,7 @@ update(s, [5, 4, 3])
 assert s == {1, 2, 3, 4, 5}
 
 [case testDictStuff]
-from typing import Dict, Any
+from typing import Dict, Any, List, Set, Tuple
 from defaultdictwrap import make_dict
 
 def f(x: int) -> int:
@@ -1035,14 +1035,22 @@ def u(x: int) -> int:
     d.update(x=x)
     return d['x']
 
+def get_content(d: Dict[int, int]) -> Tuple[List[int], List[int], List[Tuple[int, int]]]:
+    return list(d.keys()), list(d.values()), list(d.items())
+
+def get_content_set(d: Dict[int, int]) -> Tuple[Set[int], Set[int], Set[Tuple[int, int]]]:
+    return set(d.keys()), set(d.values()), set(d.items())
 [file defaultdictwrap.py]
 from typing import Dict
 from collections import defaultdict  # type: ignore
 def make_dict() -> Dict[str, int]:
     return defaultdict(int)
 
 [file driver.py]
-from native import f, g, h, u, make_dict1, make_dict2, update_dict
+from collections import OrderedDict
+from native import (
+    f, g, h, u, make_dict1, make_dict2, update_dict, get_content, get_content_set
+)
 assert f(1) == 2
 assert f(2) == 1
 assert g() == 30
@@ -1064,6 +1072,14 @@ update_dict(d, object.__dict__)
 assert d == dict(object.__dict__)
 
 assert u(10) == 10
+assert get_content({1: 2}) == ([1], [2], [(1, 2)])
+od = OrderedDict([(1, 2), (3, 4)])
+assert get_content(od) == ([1, 3], [2, 4], [(1, 2), (3, 4)])
+od.move_to_end(1)
+assert get_content(od) == ([3, 1], [4, 2], [(3, 4), (1, 2)])
+assert get_content_set({1: 2}) == ({1}, {2}, {(1, 2)})
+assert get_content_set(od) == ({1, 3}, {2, 4}, {(1, 2), (3, 4)})
+[typing fixtures/typing-full.pyi]
 
 [case testDictIterationMethodsRun]
 from typing import Dict