[mypyc] feat: true_dict_rprimitive for optimized dict fastpath usage

mypyc/analysis/ircheck.py

@@ -64,6 +64,7 @@
     RUnion,
     bytes_rprimitive,
     dict_rprimitive,
+    exact_dict_rprimitive,
     int_rprimitive,
     is_float_rprimitive,
     is_object_rprimitive,
@@ -177,6 +178,7 @@ def check_op_sources_valid(fn: FuncIR) -> list[FnError]:
     int_rprimitive.name,
     bytes_rprimitive.name,
     str_rprimitive.name,
+    exact_dict_rprimitive.name,
     dict_rprimitive.name,
     list_rprimitive.name,
     set_rprimitive.name,
@@ -197,7 +199,10 @@ def can_coerce_to(src: RType, dest: RType) -> bool:
         if isinstance(src, RPrimitive):
             # If either src or dest is a disjoint type, then they must both be.
             if src.name in disjoint_types and dest.name in disjoint_types:
-                return src.name == dest.name
+                return src.name == dest.name or (
+                    src.name in ("builtins.dict", "builtins.dict[exact]")
+                    and dest.name in ("builtins.dict", "builtins.dict[exact]")
+                )
             return src.size == dest.size
         if isinstance(src, RInstance):
             return is_object_rprimitive(dest)

mypyc/annotate.py

@@ -216,7 +216,7 @@ def function_annotations(func_ir: FuncIR, tree: MypyFile) -> dict[int, list[Anno
                         ann = "Dynamic method call."
                 elif name in op_hints:
                     ann = op_hints[name]
-                elif name in ("CPyDict_GetItem", "CPyDict_SetItem"):
+                elif name in ("CPyDict_GetItemUnsafe", "PyDict_SetItem"):
                     if (
                         isinstance(op.args[0], LoadStatic)
                         and isinstance(op.args[1], LoadLiteral)

mypyc/ir/rtypes.py

@@ -487,8 +487,15 @@ def __hash__(self) -> int:
     "builtins.list", is_unboxed=False, is_refcounted=True, may_be_immortal=False
 )
 
-# Python dict object (or an instance of a subclass of dict).
+# Python dict object.
+exact_dict_rprimitive: Final = RPrimitive(
+    "builtins.dict[exact]", is_unboxed=False, is_refcounted=True
+)
+"""A primitive for dicts that are confirmed to be actual instances of builtins.dict, not a subclass."""
+
+# An instance of a subclass of dict.
 dict_rprimitive: Final = RPrimitive("builtins.dict", is_unboxed=False, is_refcounted=True)
+"""A primitive that represents instances of builtins.dict or subclasses of dict."""
 
 # Python set object (or an instance of a subclass of set).
 set_rprimitive: Final = RPrimitive("builtins.set", is_unboxed=False, is_refcounted=True)
@@ -608,7 +615,14 @@ def is_list_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
 
 
 def is_dict_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
-    return isinstance(rtype, RPrimitive) and rtype.name == "builtins.dict"
+    return isinstance(rtype, RPrimitive) and rtype.name in (
+        "builtins.dict",
+        "builtins.dict[exact]",
+    )
+
+
+def is_exact_dict_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
+    return isinstance(rtype, RPrimitive) and rtype.name == "builtins.dict[exact]"
 
 
 def is_set_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:

mypyc/irbuild/builder.py

@@ -93,7 +93,7 @@
     bitmap_rprimitive,
     bytes_rprimitive,
     c_pyssize_t_rprimitive,
-    dict_rprimitive,
+    exact_dict_rprimitive,
     int_rprimitive,
     is_float_rprimitive,
     is_list_rprimitive,
@@ -125,7 +125,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
@@ -436,6 +436,8 @@ def add_to_non_ext_dict(
     ) -> None:
         # Add an attribute entry into the class dict of a non-extension class.
         key_unicode = self.load_str(key)
+        # must use `dict_set_item_op` instead of `exact_dict_set_item_op` because
+        # it breaks enums, and probably other stuff, if we take the fast path.
         self.primitive_op(dict_set_item_op, [non_ext.dict, key_unicode, val], line)
 
         # It's important that accessing class dictionary items from multiple threads
@@ -471,7 +473,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.
@@ -1388,10 +1390,10 @@ 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))
+        return self.add(LoadStatic(exact_dict_rprimitive, "globals", self.module_name))
 
     def load_module_attr_by_fullname(self, fullname: str, line: int) -> Value:
         module, _, name = fullname.rpartition(".")

mypyc/irbuild/classdef.py

@@ -50,7 +50,7 @@
 from mypyc.ir.rtypes import (
     RType,
     bool_rprimitive,
-    dict_rprimitive,
+    exact_dict_rprimitive,
     is_none_rprimitive,
     is_object_rprimitive,
     is_optional_type,
@@ -611,7 +611,7 @@ def setup_non_ext_dict(
         py_hasattr_op, [metaclass, builder.load_str("__prepare__")], cdef.line
     )
 
-    non_ext_dict = Register(dict_rprimitive)
+    non_ext_dict = Register(exact_dict_rprimitive)
 
     true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
     builder.add_bool_branch(has_prepare, true_block, false_block)

mypyc/irbuild/expression.py

@@ -100,7 +100,7 @@
 )
 from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
 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
@@ -186,7 +186,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/irbuild/for_helpers.py

@@ -52,6 +52,7 @@
     c_pyssize_t_rprimitive,
     int_rprimitive,
     is_dict_rprimitive,
+    is_exact_dict_rprimitive,
     is_fixed_width_rtype,
     is_immutable_rprimitive,
     is_list_rprimitive,
@@ -75,6 +76,11 @@
     dict_next_key_op,
     dict_next_value_op,
     dict_value_iter_op,
+    exact_dict_check_size_op,
+    exact_dict_iter_fast_path_op,
+    exact_dict_next_item_op,
+    exact_dict_next_key_op,
+    exact_dict_next_value_op,
 )
 from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
 from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
@@ -424,8 +430,10 @@ def make_for_loop_generator(
         # Special case "for k in <dict>".
         expr_reg = builder.accept(expr)
         target_type = builder.get_dict_key_type(expr)
-
-        for_dict = ForDictionaryKeys(builder, index, body_block, loop_exit, line, nested)
+        for_loop_cls = (
+            ForExactDictionaryKeys if is_exact_dict_rprimitive(rtyp) else ForDictionaryKeys
+        )
+        for_dict = for_loop_cls(builder, index, body_block, loop_exit, line, nested)
         for_dict.init(expr_reg, target_type)
         return for_dict
 
@@ -507,13 +515,22 @@ def make_for_loop_generator(
             for_dict_type: type[ForGenerator] | None = None
             if expr.callee.name == "keys":
                 target_type = builder.get_dict_key_type(expr.callee.expr)
-                for_dict_type = ForDictionaryKeys
+                if is_exact_dict_rprimitive(rtype):
+                    for_dict_type = ForExactDictionaryKeys
+                else:
+                    for_dict_type = ForDictionaryKeys
             elif expr.callee.name == "values":
                 target_type = builder.get_dict_value_type(expr.callee.expr)
-                for_dict_type = ForDictionaryValues
+                if is_exact_dict_rprimitive(rtype):
+                    for_dict_type = ForExactDictionaryValues
+                else:
+                    for_dict_type = ForDictionaryValues
             else:
                 target_type = builder.get_dict_item_type(expr.callee.expr)
-                for_dict_type = ForDictionaryItems
+                if is_exact_dict_rprimitive(rtype):
+                    for_dict_type = ForExactDictionaryItems
+                else:
+                    for_dict_type = ForDictionaryItems
             for_dict_gen = for_dict_type(builder, index, body_block, loop_exit, line, nested)
             for_dict_gen.init(expr_reg, target_type)
             return for_dict_gen
@@ -898,6 +915,7 @@ class ForDictionaryCommon(ForGenerator):
 
     dict_next_op: ClassVar[CFunctionDescription]
     dict_iter_op: ClassVar[CFunctionDescription]
+    dict_size_op: ClassVar[CFunctionDescription] = dict_check_size_op
 
     def need_cleanup(self) -> bool:
         # Technically, a dict subclass can raise an unrelated exception
@@ -944,7 +962,7 @@ def gen_step(self) -> None:
         line = self.line
         # Technically, we don't need a new primitive for this, but it is simpler.
         builder.call_c(
-            dict_check_size_op,
+            self.dict_size_op,
             [builder.read(self.expr_target, line), builder.read(self.size, line)],
             line,
         )
@@ -1022,6 +1040,30 @@ def begin_body(self) -> None:
             builder.assign(target, rvalue, line)
 
 
+class ForExactDictionaryKeys(ForDictionaryKeys):
+    """Generate optimized IR for a for loop over dictionary items without type checks."""
+
+    dict_next_op = exact_dict_next_key_op
+    dict_iter_op = exact_dict_iter_fast_path_op
+    dict_size_op = exact_dict_check_size_op
+
+
+class ForExactDictionaryValues(ForDictionaryValues):
+    """Generate optimized IR for a for loop over dictionary items without type checks."""
+
+    dict_next_op = exact_dict_next_value_op
+    dict_iter_op = exact_dict_iter_fast_path_op
+    dict_size_op = exact_dict_check_size_op
+
+
+class ForExactDictionaryItems(ForDictionaryItems):
+    """Generate optimized IR for a for loop over dictionary items without type checks."""
+
+    dict_next_op = exact_dict_next_item_op
+    dict_iter_op = exact_dict_iter_fast_path_op
+    dict_size_op = exact_dict_check_size_op
+
+
 class ForRange(ForGenerator):
     """Generate optimized IR for a for loop over an integer range."""
 

mypyc/irbuild/function.py

@@ -56,7 +56,7 @@
 from mypyc.ir.rtypes import (
     RInstance,
     bool_rprimitive,
-    dict_rprimitive,
+    exact_dict_rprimitive,
     int_rprimitive,
     object_rprimitive,
 )
@@ -77,8 +77,8 @@
 from mypyc.irbuild.generator import gen_generator_func, gen_generator_func_body
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.primitives.dict_ops import (
-    dict_get_method_with_none,
     dict_new_op,
+    exact_dict_get_method_with_none,
     exact_dict_set_item_op,
 )
 from mypyc.primitives.generic_ops import py_setattr_op
@@ -814,10 +814,12 @@ def generate_singledispatch_dispatch_function(
 
     arg_type = builder.builder.get_type_of_obj(arg_info.args[0], line)
     dispatch_cache = builder.builder.get_attr(
-        dispatch_func_obj, "dispatch_cache", dict_rprimitive, line
+        dispatch_func_obj, "dispatch_cache", exact_dict_rprimitive, line
     )
     call_find_impl, use_cache, call_func = BasicBlock(), BasicBlock(), BasicBlock()
-    get_result = builder.primitive_op(dict_get_method_with_none, [dispatch_cache, arg_type], line)
+    get_result = builder.primitive_op(
+        exact_dict_get_method_with_none, [dispatch_cache, arg_type], line
+    )
     is_not_none = builder.translate_is_op(get_result, builder.none_object(), "is not", line)
     impl_to_use = Register(object_rprimitive)
     builder.add_bool_branch(is_not_none, use_cache, call_find_impl)
@@ -894,8 +896,8 @@ def gen_dispatch_func_ir(
     """
     builder.enter(FuncInfo(fitem, dispatch_name))
     setup_callable_class(builder)
-    builder.fn_info.callable_class.ir.attributes["registry"] = dict_rprimitive
-    builder.fn_info.callable_class.ir.attributes["dispatch_cache"] = dict_rprimitive
+    builder.fn_info.callable_class.ir.attributes["registry"] = exact_dict_rprimitive
+    builder.fn_info.callable_class.ir.attributes["dispatch_cache"] = exact_dict_rprimitive
     builder.fn_info.callable_class.ir.has_dict = True
     builder.fn_info.callable_class.ir.needs_getseters = True
     generate_singledispatch_callable_class_ctor(builder)
@@ -958,7 +960,7 @@ def add_register_method_to_callable_class(builder: IRBuilder, fn_info: FuncInfo)
 
 
 def load_singledispatch_registry(builder: IRBuilder, dispatch_func_obj: Value, line: int) -> Value:
-    return builder.builder.get_attr(dispatch_func_obj, "registry", dict_rprimitive, line)
+    return builder.builder.get_attr(dispatch_func_obj, "registry", exact_dict_rprimitive, line)
 
 
 def singledispatch_main_func_name(orig_name: str) -> str:
@@ -1009,7 +1011,7 @@ def maybe_insert_into_registry_dict(builder: IRBuilder, fitem: FuncDef) -> None:
             loaded_type = load_type(builder, typ, None, line)
             builder.call_c(exact_dict_set_item_op, [registry, loaded_type, to_insert], line)
         dispatch_cache = builder.builder.get_attr(
-            dispatch_func_obj, "dispatch_cache", dict_rprimitive, line
+            dispatch_func_obj, "dispatch_cache", exact_dict_rprimitive, line
         )
         builder.gen_method_call(dispatch_cache, "clear", [], None, line)