[mypyc] feat: use exact dict getitem op for module dict lookup

Author: BobTheBuidler

mypyc/irbuild/builder.py

@@ -126,7 +126,7 @@
 )
 from mypyc.irbuild.util import bytes_from_str, is_constant
 from mypyc.options import CompilerOptions
-from mypyc.primitives.dict_ops import dict_get_item_op, dict_set_item_op
+from mypyc.primitives.dict_ops import dict_set_item_op, exact_dict_get_item_op
 from mypyc.primitives.generic_ops import iter_op, next_op, py_setattr_op
 from mypyc.primitives.list_ops import list_get_item_unsafe_op, list_pop_last, to_list
 from mypyc.primitives.misc_ops import check_unpack_count_op, get_module_dict_op, import_op
@@ -472,7 +472,7 @@ def get_module(self, module: str, line: int) -> Value:
         # Python 3.7 has a nice 'PyImport_GetModule' function that we can't use :(
         mod_dict = self.call_c(get_module_dict_op, [], line)
         # Get module object from modules dict.
-        return self.primitive_op(dict_get_item_op, [mod_dict, self.load_str(module)], line)
+        return self.primitive_op(exact_dict_get_item_op, [mod_dict, self.load_str(module)], line)
 
     def get_module_attr(self, module: str, attr: str, line: int) -> Value:
         """Look up an attribute of a module without storing it in the local namespace.
@@ -1410,7 +1410,7 @@ def load_global(self, expr: NameExpr) -> Value:
     def load_global_str(self, name: str, line: int) -> Value:
         _globals = self.load_globals_dict()
         reg = self.load_str(name)
-        return self.primitive_op(dict_get_item_op, [_globals, reg], line)
+        return self.primitive_op(exact_dict_get_item_op, [_globals, reg], line)
 
     def load_globals_dict(self) -> Value:
         return self.add(LoadStatic(dict_rprimitive, "globals", self.module_name))

mypyc/irbuild/expression.py

@@ -104,7 +104,7 @@
     translate_object_setattr,
 )
 from mypyc.primitives.bytes_ops import bytes_slice_op
-from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, exact_dict_set_item_op
+from mypyc.primitives.dict_ops import dict_new_op, exact_dict_get_item_op, exact_dict_set_item_op
 from mypyc.primitives.generic_ops import iter_op, name_op
 from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
 from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
@@ -190,7 +190,7 @@ def transform_name_expr(builder: IRBuilder, expr: NameExpr) -> Value:
             # instead load the module separately on each access.
             mod_dict = builder.call_c(get_module_dict_op, [], expr.line)
             obj = builder.primitive_op(
-                dict_get_item_op, [mod_dict, builder.load_str(expr.node.fullname)], expr.line
+                exact_dict_get_item_op, [mod_dict, builder.load_str(expr.node.fullname)], expr.line
             )
             return obj
         else:

mypyc/lib-rt/CPy.h

@@ -711,6 +711,9 @@ tuple_T3CIO CPyDict_NextValue(PyObject *dict_or_iter, CPyTagged offset);
 tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset);
 int CPyMapping_Check(PyObject *obj);
 
+// Unsafe dict operations (assume PyDict_CheckExact(dict) is always true)
+PyObject *CPyDict_GetItemUnsafe(PyObject *dict, PyObject *key);
+
 // Check that dictionary didn't change size during iteration.
 static inline char CPyDict_CheckSize(PyObject *dict, Py_ssize_t size) {
     if (!PyDict_CheckExact(dict)) {

mypyc/lib-rt/dict_ops.c

@@ -15,15 +15,7 @@
 // some indirections.
 PyObject *CPyDict_GetItem(PyObject *dict, PyObject *key) {
     if (PyDict_CheckExact(dict)) {
-        PyObject *res = PyDict_GetItemWithError(dict, key);
-        if (!res) {
-            if (!PyErr_Occurred()) {
-                PyErr_SetObject(PyExc_KeyError, key);
-            }
-        } else {
-            Py_INCREF(res);
-        }
-        return res;
+        return CPyDict_GetItemUnsafe(dict, key);
     } else {
         return PyObject_GetItem(dict, key);
     }
@@ -489,3 +481,22 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) {
 int CPyMapping_Check(PyObject *obj) {
     return Py_TYPE(obj)->tp_flags & Py_TPFLAGS_MAPPING;
 }
+
+// =======================
+// Unsafe dict operations
+// =======================
+
+// Unsafe: assumes dict is a true dict (PyDict_CheckExact(dict) is always true)
+
+PyObject *CPyDict_GetItemUnsafe(PyObject *dict, PyObject *key) {
+    // No type check, direct call
+    PyObject *res = PyDict_GetItemWithError(dict, key);
+    if (!res) {
+        if (!PyErr_Occurred()) {
+            PyErr_SetObject(PyExc_KeyError, key);
+        }
+    } else {
+        Py_INCREF(res);
+    }
+    return res;
+}

mypyc/primitives/dict_ops.py

@@ -81,6 +81,14 @@
 )
 
 # dict[key]
+exact_dict_get_item_op = custom_op(
+    arg_types=[dict_rprimitive, object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPyDict_GetItemUnsafe",
+    error_kind=ERR_MAGIC,
+)
+
+# dictsubclass[key]
 dict_get_item_op = method_op(
     name="__getitem__",
     arg_types=[dict_rprimitive, object_rprimitive],