[mypyc] Only construct dict setdefault collection when needed (#10668)

Closes mypyc/mypyc#830

Add dict_setdefault_spec_init_op.

Author: 97littleleaf11

mypyc/irbuild/specialize.py

@@ -14,17 +14,22 @@
 
 from typing import Callable, Optional, Dict, Tuple, List
 
-from mypy.nodes import CallExpr, RefExpr, MemberExpr, TupleExpr, GeneratorExpr, ARG_POS
+from mypy.nodes import (
+    CallExpr, RefExpr, MemberExpr, NameExpr, TupleExpr, GeneratorExpr,
+    ListExpr, DictExpr, ARG_POS
+)
 from mypy.types import AnyType, TypeOfAny
 
 from mypyc.ir.ops import (
     Value, Register, BasicBlock, Integer, RaiseStandardError, Unreachable
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive, is_dict_rprimitive
+    bool_rprimitive, is_dict_rprimitive, c_int_rprimitive
+)
+from mypyc.primitives.dict_ops import (
+    dict_keys_op, dict_values_op, dict_items_op, dict_setdefault_spec_init_op
 )
-from mypyc.primitives.dict_ops import dict_keys_op, dict_values_op, dict_items_op
 from mypyc.primitives.list_ops import new_list_set_item_op
 from mypyc.primitives.tuple_ops import new_tuple_set_item_op
 from mypyc.irbuild.builder import IRBuilder
@@ -313,3 +318,41 @@ def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
         if irs is not None:
             return builder.builder.isinstance_helper(builder.accept(expr.args[0]), irs, expr.line)
     return None
+
+
+@specialize_function('setdefault', dict_rprimitive)
+def translate_dict_setdefault(
+        builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
+    if (len(expr.args) == 2
+            and expr.arg_kinds == [ARG_POS, ARG_POS]
+            and isinstance(callee, MemberExpr)):
+        # Special case for dict.setdefault which would only construct default empty
+        # collection when needed. The dict_setdefault_spec_init_op checks whether
+        # the dict contains the key and would construct the empty collection only once.
+        # For example, this specializer works for the following cases:
+        #     d.setdefault(key, set()).add(value)
+        #     d.setdefault(key, []).append(value)
+        #     d.setdefault(key, {})[inner_key] = inner_val
+        arg = expr.args[1]
+        if isinstance(arg, ListExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(1, c_int_rprimitive, expr.line)
+        elif isinstance(arg, DictExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(2, c_int_rprimitive, expr.line)
+        elif (isinstance(arg, CallExpr) and isinstance(arg.callee, NameExpr)
+                and arg.callee.fullname == 'builtins.set'):
+            if len(arg.args):
+                return None
+            data_type = Integer(3, c_int_rprimitive, expr.line)
+        else:
+            return None
+
+        callee_dict = builder.accept(callee.expr)
+        key_val = builder.accept(expr.args[0])
+        return builder.call_c(dict_setdefault_spec_init_op,
+                              [callee_dict, key_val, data_type],
+                              expr.line)
+    return None

mypyc/lib-rt/CPy.h

@@ -344,6 +344,7 @@ PyObject *CPyDict_Get(PyObject *dict, PyObject *key, PyObject *fallback);
 PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key);
 PyObject *CPyDict_SetDefault(PyObject *dict, PyObject *key, PyObject *value);
 PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key);
+PyObject *CPyDict_SetDefaultWithEmptyDatatype(PyObject *dict, PyObject *key, int data_type);
 PyObject *CPyDict_Build(Py_ssize_t size, ...);
 int CPyDict_Update(PyObject *dict, PyObject *stuff);
 int CPyDict_UpdateInDisplay(PyObject *dict, PyObject *stuff);

mypyc/lib-rt/dict_ops.c

@@ -80,6 +80,35 @@ PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key) {
     return CPyDict_SetDefault(dict, key, Py_None);
 }
 
+PyObject *CPyDict_SetDefaultWithEmptyDatatype(PyObject *dict, PyObject *key,
+                                              int data_type) {
+    PyObject *res = CPyDict_GetItem(dict, key);
+    if (!res) {
+        // CPyDict_GetItem() would generates an PyExc_KeyError
+        // when key is not found.
+        PyErr_Clear();
+
+        PyObject *new_obj;
+        if (data_type == 1) {
+            new_obj = PyList_New(0);
+        } else if (data_type == 2) {
+            new_obj = PyDict_New();
+        } else if (data_type == 3) {
+            new_obj = PySet_New(NULL);
+        } else {
+            return NULL;
+        }
+
+        if (CPyDict_SetItem(dict, key, new_obj) == -1) {
+            return NULL;
+        } else {
+            return new_obj;
+        }
+    } else {
+        return res;
+    }
+}
+
 int CPyDict_SetItem(PyObject *dict, PyObject *key, PyObject *value) {
     if (PyDict_CheckExact(dict)) {
         return PyDict_SetItem(dict, key, value);

mypyc/primitives/dict_ops.py

@@ -119,7 +119,7 @@
     error_kind=ERR_MAGIC)
 
 # dict.setdefault(key, default)
-method_op(
+dict_setdefault_op = method_op(
     name='setdefault',
     arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,
@@ -135,6 +135,16 @@
     is_borrowed=True,
     error_kind=ERR_MAGIC)
 
+# dict.setdefault(key, empty tuple/list/set)
+# The third argument marks the data type of the second argument.
+#     1: list    2: dict    3: set
+# Other number would lead to an error.
+dict_setdefault_spec_init_op = custom_op(
+    arg_types=[dict_rprimitive, object_rprimitive, c_int_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name='CPyDict_SetDefaultWithEmptyDatatype',
+    error_kind=ERR_MAGIC)
+
 # dict.keys()
 method_op(
     name='keys',

mypyc/test-data/fixtures/ir.py

@@ -177,7 +177,7 @@ def values(self) -> Iterable[V]: pass
     def items(self) -> Iterable[Tuple[K, V]]: pass
     def clear(self) -> None: pass
     def copy(self) -> Dict[K, V]: pass
-    def setdefault(self, k: K, v: Optional[V] = None) -> Optional[V]: pass
+    def setdefault(self, key: K, val: V = ...) -> V: pass
 
 class set(Generic[T]):
     def __init__(self, i: Optional[Iterable[T]] = None) -> None: pass

mypyc/test-data/irbuild-dict.test

@@ -340,6 +340,24 @@ L0:
 from typing import Dict
 def f(d: Dict[object, object]) -> object:
     return d.setdefault('a', 'b')
+
+def f2(d: Dict[object, object], flag: bool) -> object:
+    if flag:
+        return d.setdefault('a', set())
+    else:
+        return d.setdefault('a', set('b'))
+
+def f3(d: Dict[object, object], flag: bool) -> object:
+    if flag:
+        return d.setdefault('a', [])
+    else:
+        return d.setdefault('a', [1])
+
+def f4(d: Dict[object, object], flag: bool) -> object:
+    if flag:
+        return d.setdefault('a', {})
+    else:
+        return d.setdefault('a', {'c': 1})
 [out]
 def f(d):
     d :: dict
@@ -350,3 +368,81 @@ L0:
     r1 = 'b'
     r2 = CPyDict_SetDefault(d, r0, r1)
     return r2
+def f2(d, flag):
+    d :: dict
+    flag :: bool
+    r0 :: str
+    r1 :: object
+    r2, r3 :: str
+    r4 :: set
+    r5, r6 :: object
+L0:
+    if flag goto L1 else goto L2 :: bool
+L1:
+    r0 = 'a'
+    r1 = CPyDict_SetDefaultWithEmptyDatatype(d, r0, 3)
+    return r1
+L2:
+    r2 = 'a'
+    r3 = 'b'
+    r4 = PySet_New(r3)
+    r5 = CPyDict_SetDefault(d, r2, r4)
+    return r5
+L3:
+    r6 = box(None, 1)
+    return r6
+def f3(d, flag):
+    d :: dict
+    flag :: bool
+    r0 :: str
+    r1 :: object
+    r2 :: str
+    r3 :: list
+    r4 :: object
+    r5, r6 :: ptr
+    r7, r8 :: object
+L0:
+    if flag goto L1 else goto L2 :: bool
+L1:
+    r0 = 'a'
+    r1 = CPyDict_SetDefaultWithEmptyDatatype(d, r0, 1)
+    return r1
+L2:
+    r2 = 'a'
+    r3 = PyList_New(1)
+    r4 = box(short_int, 2)
+    r5 = get_element_ptr r3 ob_item :: PyListObject
+    r6 = load_mem r5 :: ptr*
+    set_mem r6, r4 :: builtins.object*
+    keep_alive r3
+    r7 = CPyDict_SetDefault(d, r2, r3)
+    return r7
+L3:
+    r8 = box(None, 1)
+    return r8
+def f4(d, flag):
+    d :: dict
+    flag :: bool
+    r0 :: str
+    r1 :: object
+    r2, r3 :: str
+    r4 :: object
+    r5 :: dict
+    r6, r7 :: object
+L0:
+    if flag goto L1 else goto L2 :: bool
+L1:
+    r0 = 'a'
+    r1 = CPyDict_SetDefaultWithEmptyDatatype(d, r0, 2)
+    return r1
+L2:
+    r2 = 'a'
+    r3 = 'c'
+    r4 = box(short_int, 2)
+    r5 = CPyDict_Build(1, r3, r4)
+    r6 = CPyDict_SetDefault(d, r2, r5)
+    return r6
+L3:
+    r7 = box(None, 1)
+    return r7
+

mypyc/test-data/run-dicts.test

@@ -196,7 +196,7 @@ else:
 
 [case testDictMethods]
 from collections import defaultdict
-from typing import Dict, Optional
+from typing import Dict, Optional, List, Set
 
 def test_dict_clear() -> None:
     d = {'a': 1, 'b': 2}
@@ -232,7 +232,7 @@ class MyDict(dict):
             return super().setdefault(k, v) + 10
 
 def test_dict_setdefault() -> None:
-    d: Dict[str, int] = {'a': 1, 'b': 2}
+    d: Dict[str, Optional[int]] = {'a': 1, 'b': 2}
     assert d.setdefault('a', 2) == 1
     assert d.setdefault('b', 2) == 2
     assert d.setdefault('c', 3) == 3
@@ -254,6 +254,32 @@ def test_dict_subclass_setdefault() -> None:
     assert d.setdefault('e') == None
     assert d.setdefault('e', 100) == 110
 
+def test_dict_empty_collection_setdefault() -> None:
+    d1: Dict[str, List[int]] = {'a': [1, 2, 3]}
+    assert d1.setdefault('a', []) == [1, 2, 3]
+    assert d1.setdefault('b', []) == []
+    assert 'b' in d1
+    d1.setdefault('b', []).append(3)
+    assert d1['b'] == [3]
+    assert d1.setdefault('c', [1]) == [1]
+
+    d2: Dict[str, Dict[str, int]] = {'a': {'a': 1}}
+    assert d2.setdefault('a', {}) == {'a': 1}
+    assert d2.setdefault('b', {}) == {}
+    assert 'b' in d2
+    d2.setdefault('b', {})['aa'] = 2
+    d2.setdefault('b', {})['bb'] = 3
+    assert d2['b'] == {'aa': 2, 'bb': 3}
+    assert d2.setdefault('c', {'cc': 1}) == {'cc': 1}
+
+    d3: Dict[str, Set[str]] = {'a': set('a')}
+    assert d3.setdefault('a', set()) == {'a'}
+    assert d3.setdefault('b', set()) == set()
+    d3.setdefault('b', set()).add('b')
+    d3.setdefault('b', set()).add('c')
+    assert d3['b'] == {'b', 'c'}
+    assert d3.setdefault('c', set('d')) == {'d'}
+
 [case testDictToBool]
 from typing import Dict, List