Repository: python/mypy
Pull request: [mypyc] Only construct default collection when needed
### Description

Closes https://github.com/mypyc/mypyc/issues/830

* Add `dict_setdefault_spec_init_op`

## Test Plan

* Add new run tests for empty collection as defaults
* Add irbuild tests


Problem statement:
Optimize setdefault to only construct default collection when needed
Right now this code always constructs an empty set, even if it's not needed:

```py
d = {}
...
d.setdefault(key, set()).add(value)
```

If the key already has a value, the newly constructed `set()` is discarded immediately.

We could make this faster by detecting certain default values and constructing them lazily. These default values should be sufficient:

* `[]`
* `{}`
* `set()`

Implementation sketch:
* Add a new C primitive for `dict.setdefault` that is used only if the default argument is one of the above expressions (exactly).
* Instead of a `PyObject *` default object, pass an integer with 3 possible values correponding to the different types (`list`, `dict` and `set`).
* In the primitive, try to look up the key from the dictionary. If it exists, return the value without constructing the container. Otherwise, construct the container and add it to the dict.

Another option would be to inline the implementation of `setdefault` in IR. If the arguments would be evaluated only as needed, this would also give us the desired functionality. This may be too complicated, though.

This is not high priority but could still be worth it, since it's common to have an empty container as the default.

Hints from issue discussion:

Commit message: Add tests

Patch to review:
diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py
index ef2986cda0b1..cf524fe96fd8 100644
--- a/mypyc/irbuild/specialize.py
+++ b/mypyc/irbuild/specialize.py
@@ -14,7 +14,10 @@
 
 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 (
@@ -22,9 +25,11 @@
 )
 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
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index f1a342527b11..b50802f2d6af 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/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);
diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c
index 2de5f4a9e7a4..be9bf4432192 100644
--- a/mypyc/lib-rt/dict_ops.c
+++ b/mypyc/lib-rt/dict_ops.c
@@ -80,6 +80,31 @@ 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) {
+        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);
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index e176ffcac316..c73d9e615d19 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/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,13 @@
     is_borrowed=True,
     error_kind=ERR_MAGIC)
 
+# dict.setdefault(key, empty tuple/list/set)
+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',
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index 9a64280c5372..91c17e727cf7 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/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
diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index da08ed79d5bd..bfded5d6309f 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/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
+
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index b25b5237e80f..41675e7fcc91 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/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
 


Code review comments:
diff_hunk: @@ -80,6 +80,31 @@ 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) {
+        PyErr_Clear();
path: mypyc/lib-rt/dict_ops.c
text: @user1:
Do you have this because the function may return `NULL` without raising an exception in case `data_type` is unsupported? It would be cleaner to instead raise `RuntimeError` or similar instead of just returning `NULL`. At least it would be better to have this just before `return NULL;` to make it clearer why we are doing this.