[mypyc] feat: exact_dict_set_item_op (#19657)

This PR implements a `exact_dict_set_item_op` custom_op in preparation
for the addition of exact_dict_rprimitive. We don't actually need to
implement exact_dict_rprimitive before implementing this op as long as
we accept that it will not be used for CallExpr specialization. This is
okay for us now, as it is already addressed in the larger PR.

Author: BobTheBuidler

mypyc/irbuild/classdef.py

@@ -66,7 +66,7 @@
 )
 from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator
-from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
+from mypyc.primitives.dict_ops import dict_new_op, exact_dict_set_item_op
 from mypyc.primitives.generic_ops import (
     iter_op,
     next_op,
@@ -271,8 +271,8 @@ def finalize(self, ir: ClassIR) -> None:
         )
 
         # Add the non-extension class to the dict
-        self.builder.primitive_op(
-            dict_set_item_op,
+        self.builder.call_c(
+            exact_dict_set_item_op,
             [
                 self.builder.load_globals_dict(),
                 self.builder.load_str(self.cdef.name),
@@ -487,8 +487,10 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
     builder.add(InitStatic(tp, cdef.name, builder.module_name, NAMESPACE_TYPE))
 
     # Add it to the dict
-    builder.primitive_op(
-        dict_set_item_op, [builder.load_globals_dict(), builder.load_str(cdef.name), tp], cdef.line
+    builder.call_c(
+        exact_dict_set_item_op,
+        [builder.load_globals_dict(), builder.load_str(cdef.name), tp],
+        cdef.line,
     )
 
     return tp
@@ -672,7 +674,7 @@ def add_non_ext_class_attr_ann(
             typ = builder.add(LoadAddress(type_object_op.type, type_object_op.src, stmt.line))
 
     key = builder.load_str(lvalue.name)
-    builder.primitive_op(dict_set_item_op, [non_ext.anns, key, typ], stmt.line)
+    builder.call_c(exact_dict_set_item_op, [non_ext.anns, key, typ], stmt.line)
 
 
 def add_non_ext_class_attr(

mypyc/irbuild/expression.py

@@ -97,7 +97,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, dict_set_item_op
+from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, exact_dict_set_item_op
 from mypyc.primitives.generic_ops import iter_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
@@ -1030,7 +1030,7 @@ def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehe
     def gen_inner_stmts() -> None:
         k = builder.accept(o.key)
         v = builder.accept(o.value)
-        builder.primitive_op(dict_set_item_op, [builder.read(d), k, v], o.line)
+        builder.call_c(exact_dict_set_item_op, [builder.read(d), k, v], o.line)
 
     comprehension_helper(builder, loop_params, gen_inner_stmts, o.line)
     return builder.read(d)

mypyc/irbuild/function.py

@@ -76,7 +76,11 @@
 )
 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, dict_set_item_op
+from mypyc.primitives.dict_ops import (
+    dict_get_method_with_none,
+    dict_new_op,
+    exact_dict_set_item_op,
+)
 from mypyc.primitives.generic_ops import py_setattr_op
 from mypyc.primitives.misc_ops import register_function
 from mypyc.primitives.registry import builtin_names
@@ -123,8 +127,8 @@ def transform_decorator(builder: IRBuilder, dec: Decorator) -> None:
 
     if decorated_func is not None:
         # Set the callable object representing the decorated function as a global.
-        builder.primitive_op(
-            dict_set_item_op,
+        builder.call_c(
+            exact_dict_set_item_op,
             [builder.load_globals_dict(), builder.load_str(dec.func.name), decorated_func],
             decorated_func.line,
         )
@@ -826,7 +830,7 @@ def generate_singledispatch_dispatch_function(
     find_impl = builder.load_module_attr_by_fullname("functools._find_impl", line)
     registry = load_singledispatch_registry(builder, dispatch_func_obj, line)
     uncached_impl = builder.py_call(find_impl, [arg_type, registry], line)
-    builder.primitive_op(dict_set_item_op, [dispatch_cache, arg_type, uncached_impl], line)
+    builder.call_c(exact_dict_set_item_op, [dispatch_cache, arg_type, uncached_impl], line)
     builder.assign(impl_to_use, uncached_impl, line)
     builder.goto(call_func)
 
@@ -1003,7 +1007,7 @@ def maybe_insert_into_registry_dict(builder: IRBuilder, fitem: FuncDef) -> None:
         registry = load_singledispatch_registry(builder, dispatch_func_obj, line)
         for typ in types:
             loaded_type = load_type(builder, typ, None, line)
-            builder.primitive_op(dict_set_item_op, [registry, loaded_type, to_insert], 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
         )

mypyc/primitives/dict_ops.py

@@ -98,6 +98,15 @@
     error_kind=ERR_NEG_INT,
 )
 
+# dict[key] = value (exact dict only, no subclasses)
+# NOTE: this is currently for internal use only, and not used for CallExpr specialization
+exact_dict_set_item_op = custom_op(
+    arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name="PyDict_SetItem",
+    error_kind=ERR_NEG_INT,
+)
+
 # key in dict
 binary_op(
     name="in",

mypyc/test-data/irbuild-basic.test

@@ -1993,7 +1993,7 @@ L6:
     r13 = CPyTagged_Multiply(x, x)
     r14 = box(int, x)
     r15 = box(int, r13)
-    r16 = CPyDict_SetItem(r0, r14, r15)
+    r16 = PyDict_SetItem(r0, r14, r15)
     r17 = r16 >= 0 :: signed
 L7:
     r18 = r6 + 1
@@ -2620,7 +2620,7 @@ L0:
     d = r14
     r15 = __main__.globals :: static
     r16 = 'd'
-    r17 = CPyDict_SetItem(r15, r16, r14)
+    r17 = PyDict_SetItem(r15, r16, r14)
     r18 = r17 >= 0 :: signed
     r19 = 'c'
     r20 = builtins :: module
@@ -2693,7 +2693,7 @@ L2:
     keep_alive r17
     r24 = __main__.globals :: static
     r25 = 'c'
-    r26 = CPyDict_SetItem(r24, r25, r23)
+    r26 = PyDict_SetItem(r24, r25, r23)
     r27 = r26 >= 0 :: signed
     return 1
 

mypyc/test-data/irbuild-classes.test

@@ -303,7 +303,7 @@ L2:
     __main__.C = r27 :: type
     r33 = __main__.globals :: static
     r34 = 'C'
-    r35 = CPyDict_SetItem(r33, r34, r27)
+    r35 = PyDict_SetItem(r33, r34, r27)
     r36 = r35 >= 0 :: signed
     r37 = <error> :: object
     r38 = '__main__'
@@ -316,7 +316,7 @@ L2:
     __main__.S = r40 :: type
     r45 = __main__.globals :: static
     r46 = 'S'
-    r47 = CPyDict_SetItem(r45, r46, r40)
+    r47 = PyDict_SetItem(r45, r46, r40)
     r48 = r47 >= 0 :: signed
     r49 = __main__.C :: type
     r50 = __main__.S :: type
@@ -340,7 +340,7 @@ L2:
     __main__.D = r61 :: type
     r68 = __main__.globals :: static
     r69 = 'D'
-    r70 = CPyDict_SetItem(r68, r69, r61)
+    r70 = PyDict_SetItem(r68, r69, r61)
     r71 = r70 >= 0 :: signed
     return 1
 

mypyc/test-data/irbuild-generics.test

@@ -351,7 +351,7 @@ L17:
     k = r49
     v = r50
     r51 = box(int, v)
-    r52 = CPyDict_SetItem(r40, k, r51)
+    r52 = PyDict_SetItem(r40, k, r51)
     r53 = r52 >= 0 :: signed
 L18:
     r54 = CPyDict_CheckSize(m, r42)
@@ -493,7 +493,7 @@ L17:
     r49 = cast(union[int, str], r47)
     k = r48
     v = r49
-    r50 = CPyDict_SetItem(r39, k, v)
+    r50 = PyDict_SetItem(r39, k, v)
     r51 = r50 >= 0 :: signed
 L18:
     r52 = CPyDict_CheckSize(m, r41)
@@ -630,7 +630,7 @@ L17:
     r47 = cast(str, r45)
     k = r47
     v = r46
-    r48 = CPyDict_SetItem(r38, k, v)
+    r48 = PyDict_SetItem(r38, k, v)
     r49 = r48 >= 0 :: signed
 L18:
     r50 = CPyDict_CheckSize(t, r40)
@@ -742,7 +742,7 @@ L6:
     r17 = cast(str, r15)
     k = r17
     v = r16
-    r18 = CPyDict_SetItem(r8, k, v)
+    r18 = PyDict_SetItem(r8, k, v)
     r19 = r18 >= 0 :: signed
 L7:
     r20 = CPyDict_CheckSize(kwargs, r10)

mypyc/test-data/irbuild-singledispatch.test

@@ -76,7 +76,7 @@ L2:
     r12 = load_address r11
     r13 = PyObject_Vectorcall(r9, r12, 2, 0)
     keep_alive r1, r10
-    r14 = CPyDict_SetItem(r2, r1, r13)
+    r14 = PyDict_SetItem(r2, r1, r13)
     r15 = r14 >= 0 :: signed
     r6 = r13
 L3:
@@ -214,7 +214,7 @@ L2:
     r12 = load_address r11
     r13 = PyObject_Vectorcall(r9, r12, 2, 0)
     keep_alive r1, r10
-    r14 = CPyDict_SetItem(r2, r1, r13)
+    r14 = PyDict_SetItem(r2, r1, r13)
     r15 = r14 >= 0 :: signed
     r6 = r13
 L3: