From 98eee40c8e6b111361b3e2d69059b445852fbda3 Mon Sep 17 00:00:00 2001
From: vsakkas <vasileios.sakkas96@gmail.com>
Date: Wed, 2 Dec 2020 19:33:06 +0200
Subject: [PATCH] [mypyc] Implement dict clear primitive (#9724)

Implements dict clear primitive for improved performance.

Related ticket: mypyc/mypyc#644
---
 mypyc/lib-rt/CPy.h                |  1 +
 mypyc/lib-rt/dict_ops.c           | 12 ++++++++++++
 mypyc/primitives/dict_ops.py      |  8 ++++++++
 mypyc/test-data/irbuild-dict.test | 12 ++++++++++++
 mypyc/test-data/run-dicts.test    | 14 ++++++++++++++
 5 files changed, 47 insertions(+)

diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index 4eab87e4c011..1abc304677a4 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/mypyc/lib-rt/CPy.h
@@ -348,6 +348,7 @@ PyObject *CPyDict_ItemsView(PyObject *dict);
 PyObject *CPyDict_Keys(PyObject *dict);
 PyObject *CPyDict_Values(PyObject *dict);
 PyObject *CPyDict_Items(PyObject *dict);
+char CPyDict_Clear(PyObject *dict);
 PyObject *CPyDict_GetKeysIter(PyObject *dict);
 PyObject *CPyDict_GetItemsIter(PyObject *dict);
 PyObject *CPyDict_GetValuesIter(PyObject *dict);
diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c
index 52ccc2c94b77..fca096468ddd 100644
--- a/mypyc/lib-rt/dict_ops.c
+++ b/mypyc/lib-rt/dict_ops.c
@@ -226,6 +226,18 @@ PyObject *CPyDict_Items(PyObject *dict) {
     return list;
 }
 
+char CPyDict_Clear(PyObject *dict) {
+    if (PyDict_CheckExact(dict)) {
+        PyDict_Clear(dict);
+    } else {
+        PyObject *res = PyObject_CallMethod(dict, "clear", NULL);
+        if (res == NULL) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
 PyObject *CPyDict_GetKeysIter(PyObject *dict) {
     if (PyDict_CheckExact(dict)) {
         // Return dict itself to indicate we can use fast path instead.
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index 2838a7c06b17..e0278182d5fe 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -142,6 +142,14 @@
     c_function_name='CPyDict_ItemsView',
     error_kind=ERR_MAGIC)
 
+# dict.clear()
+method_op(
+    name='clear',
+    arg_types=[dict_rprimitive],
+    return_type=bit_rprimitive,
+    c_function_name='CPyDict_Clear',
+    error_kind=ERR_FALSE)
+
 # list(dict.keys())
 dict_keys_op = custom_op(
     arg_types=[dict_rprimitive],
diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index 37bbd09d1cef..3ba24c448972 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/mypyc/test-data/irbuild-dict.test
@@ -312,3 +312,15 @@ L0:
     r0 = load_address PyDict_Type
     x = r0
     return 1
+
+[case testDictClear]
+from typing import Dict
+def f(d: Dict[int, int]) -> None:
+    return d.clear()
+[out]
+def f(d):
+    d :: dict
+    r0 :: bit
+L0:
+    r0 = CPyDict_Clear(d)
+    return 1
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index cac68b9af060..1955d514d43a 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/mypyc/test-data/run-dicts.test
@@ -91,6 +91,7 @@ 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]
@@ -192,3 +193,16 @@ else:
 2
 4
 2
+
+[case testDictMethods]
+from collections import defaultdict
+from typing import Dict
+
+def test_dict_clear() -> None:
+    d = {'a': 1, 'b': 2}
+    d.clear()
+    assert d == {}
+    dd: Dict[str, int] = defaultdict(int)
+    dd['a'] = 1
+    dd.clear()
+    assert dd == {}