[mypyc] feat: exact_dict_set_item_op

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,8 @@ 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 +672,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,7 @@
 )
 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 +123,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 +826,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 +1003,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

@@ -90,6 +90,15 @@
 )
 
 # dict[key] = value
+# 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,
+)
+
+# dictorsubclass[key] = value
 dict_set_item_op = method_op(
     name="__setitem__",
     arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],

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-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: