Skip to content

Commit dd6cded

Browse files
committed
feat: true dict rprimitive
1 parent 2e5d7ee commit dd6cded

26 files changed

+973
-463
lines changed
20.2 KB
Binary file not shown.

mypyc/ir/rtypes.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -487,8 +487,13 @@ def __hash__(self) -> int:
487487
"builtins.list", is_unboxed=False, is_refcounted=True, may_be_immortal=False
488488
)
489489

490-
# Python dict object (or an instance of a subclass of dict).
490+
# Python dict object.
491+
true_dict_rprimitive: Final = RPrimitive("builtins.dict[confirmed]", is_unboxed=False, is_refcounted=True)
492+
"""A primitive for dicts that are confirmed to be actual instances of builtins.dict, not a subclass."""
493+
494+
# An instance of a subclass of dict.
491495
dict_rprimitive: Final = RPrimitive("builtins.dict", is_unboxed=False, is_refcounted=True)
496+
"""A primitive that represents instances of builtins.dict or subclasses of dict."""
492497

493498
# Python set object (or an instance of a subclass of set).
494499
set_rprimitive: Final = RPrimitive("builtins.set", is_unboxed=False, is_refcounted=True)
@@ -599,7 +604,11 @@ def is_list_rprimitive(rtype: RType) -> bool:
599604

600605

601606
def is_dict_rprimitive(rtype: RType) -> bool:
602-
return isinstance(rtype, RPrimitive) and rtype.name == "builtins.dict"
607+
return isinstance(rtype, RPrimitive) and rtype.name in ("builtins.dict", "builtins.dict[confirmed]")
608+
609+
610+
def is_true_dict_rprimitive(rtype: RType) -> bool:
611+
return isinstance(rtype, RPrimitive) and rtype.name == "builtins.dict[confirmed]"
603612

604613

605614
def is_set_rprimitive(rtype: RType) -> bool:

mypyc/irbuild/builder.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,6 @@
9292
RUnion,
9393
bitmap_rprimitive,
9494
c_pyssize_t_rprimitive,
95-
dict_rprimitive,
9695
int_rprimitive,
9796
is_float_rprimitive,
9897
is_list_rprimitive,
@@ -103,6 +102,7 @@
103102
none_rprimitive,
104103
object_rprimitive,
105104
str_rprimitive,
105+
true_dict_rprimitive,
106106
)
107107
from mypyc.irbuild.context import FuncInfo, ImplicitClass
108108
from mypyc.irbuild.ll_builder import LowLevelIRBuilder
@@ -124,7 +124,7 @@
124124
)
125125
from mypyc.irbuild.util import bytes_from_str, is_constant
126126
from mypyc.options import CompilerOptions
127-
from mypyc.primitives.dict_ops import dict_get_item_op, dict_set_item_op
127+
from mypyc.primitives.dict_ops import true_dict_get_item_op, true_dict_set_item_op
128128
from mypyc.primitives.generic_ops import iter_op, next_op, py_setattr_op
129129
from mypyc.primitives.list_ops import list_get_item_unsafe_op, list_pop_last, to_list
130130
from mypyc.primitives.misc_ops import check_unpack_count_op, get_module_dict_op, import_op
@@ -431,7 +431,7 @@ def add_to_non_ext_dict(
431431
) -> None:
432432
# Add an attribute entry into the class dict of a non-extension class.
433433
key_unicode = self.load_str(key)
434-
self.primitive_op(dict_set_item_op, [non_ext.dict, key_unicode, val], line)
434+
self.primitive_op(true_dict_set_item_op, [non_ext.dict, key_unicode, val], line)
435435

436436
def gen_import(self, id: str, line: int) -> None:
437437
self.imports[id] = None
@@ -462,7 +462,7 @@ def get_module(self, module: str, line: int) -> Value:
462462
# Python 3.7 has a nice 'PyImport_GetModule' function that we can't use :(
463463
mod_dict = self.call_c(get_module_dict_op, [], line)
464464
# Get module object from modules dict.
465-
return self.primitive_op(dict_get_item_op, [mod_dict, self.load_str(module)], line)
465+
return self.primitive_op(true_dict_get_item_op, [mod_dict, self.load_str(module)], line)
466466

467467
def get_module_attr(self, module: str, attr: str, line: int) -> Value:
468468
"""Look up an attribute of a module without storing it in the local namespace.
@@ -1370,10 +1370,10 @@ def load_global(self, expr: NameExpr) -> Value:
13701370
def load_global_str(self, name: str, line: int) -> Value:
13711371
_globals = self.load_globals_dict()
13721372
reg = self.load_str(name)
1373-
return self.primitive_op(dict_get_item_op, [_globals, reg], line)
1373+
return self.primitive_op(true_dict_get_item_op, [_globals, reg], line)
13741374

13751375
def load_globals_dict(self) -> Value:
1376-
return self.add(LoadStatic(dict_rprimitive, "globals", self.module_name))
1376+
return self.add(LoadStatic(true_dict_rprimitive, "globals", self.module_name))
13771377

13781378
def load_module_attr_by_fullname(self, fullname: str, line: int) -> Value:
13791379
module, _, name = fullname.rpartition(".")

mypyc/irbuild/classdef.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,11 @@
5050
from mypyc.ir.rtypes import (
5151
RType,
5252
bool_rprimitive,
53-
dict_rprimitive,
5453
is_none_rprimitive,
5554
is_object_rprimitive,
5655
is_optional_type,
5756
object_rprimitive,
57+
true_dict_rprimitive,
5858
)
5959
from mypyc.irbuild.builder import IRBuilder, create_type_params
6060
from mypyc.irbuild.function import (
@@ -66,7 +66,7 @@
6666
)
6767
from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
6868
from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator
69-
from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
69+
from mypyc.primitives.dict_ops import dict_new_op, true_dict_set_item_op
7070
from mypyc.primitives.generic_ops import (
7171
iter_op,
7272
next_op,
@@ -269,7 +269,7 @@ def finalize(self, ir: ClassIR) -> None:
269269

270270
# Add the non-extension class to the dict
271271
self.builder.primitive_op(
272-
dict_set_item_op,
272+
true_dict_set_item_op,
273273
[
274274
self.builder.load_globals_dict(),
275275
self.builder.load_str(self.cdef.name),
@@ -480,7 +480,7 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
480480

481481
# Add it to the dict
482482
builder.primitive_op(
483-
dict_set_item_op, [builder.load_globals_dict(), builder.load_str(cdef.name), tp], cdef.line
483+
true_dict_set_item_op, [builder.load_globals_dict(), builder.load_str(cdef.name), tp], cdef.line
484484
)
485485

486486
return tp
@@ -601,7 +601,7 @@ def setup_non_ext_dict(
601601
py_hasattr_op, [metaclass, builder.load_str("__prepare__")], cdef.line
602602
)
603603

604-
non_ext_dict = Register(dict_rprimitive)
604+
non_ext_dict = Register(true_dict_rprimitive)
605605

606606
true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
607607
builder.add_bool_branch(has_prepare, true_block, false_block)
@@ -664,7 +664,7 @@ def add_non_ext_class_attr_ann(
664664
typ = builder.add(LoadAddress(type_object_op.type, type_object_op.src, stmt.line))
665665

666666
key = builder.load_str(lvalue.name)
667-
builder.primitive_op(dict_set_item_op, [non_ext.anns, key, typ], stmt.line)
667+
builder.primitive_op(true_dict_set_item_op, [non_ext.anns, key, typ], stmt.line)
668668

669669

670670
def add_non_ext_class_attr(

mypyc/irbuild/expression.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@
9797
)
9898
from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
9999
from mypyc.primitives.bytes_ops import bytes_slice_op
100-
from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, dict_set_item_op
100+
from mypyc.primitives.dict_ops import dict_new_op, true_dict_get_item_op, true_dict_set_item_op
101101
from mypyc.primitives.generic_ops import iter_op
102102
from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
103103
from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
@@ -183,7 +183,7 @@ def transform_name_expr(builder: IRBuilder, expr: NameExpr) -> Value:
183183
# instead load the module separately on each access.
184184
mod_dict = builder.call_c(get_module_dict_op, [], expr.line)
185185
obj = builder.primitive_op(
186-
dict_get_item_op, [mod_dict, builder.load_str(expr.node.fullname)], expr.line
186+
true_dict_get_item_op, [mod_dict, builder.load_str(expr.node.fullname)], expr.line
187187
)
188188
return obj
189189
else:
@@ -1030,7 +1030,7 @@ def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehe
10301030
def gen_inner_stmts() -> None:
10311031
k = builder.accept(o.key)
10321032
v = builder.accept(o.value)
1033-
builder.primitive_op(dict_set_item_op, [builder.read(d), k, v], o.line)
1033+
builder.primitive_op(true_dict_set_item_op, [builder.read(d), k, v], o.line)
10341034

10351035
comprehension_helper(builder, loop_params, gen_inner_stmts, o.line)
10361036
return builder.read(d)

mypyc/irbuild/for_helpers.py

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
is_sequence_rprimitive,
5353
is_short_int_rprimitive,
5454
is_str_rprimitive,
55+
is_true_dict_rprimitive,
5556
is_tuple_rprimitive,
5657
object_pointer_rprimitive,
5758
object_rprimitive,
@@ -69,6 +70,11 @@
6970
dict_next_key_op,
7071
dict_next_value_op,
7172
dict_value_iter_op,
73+
true_dict_check_size_op,
74+
true_dict_iter_fast_path_op,
75+
true_dict_next_item_op,
76+
true_dict_next_key_op,
77+
true_dict_next_value_op,
7278
)
7379
from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
7480
from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
@@ -415,8 +421,8 @@ def make_for_loop_generator(
415421
# Special case "for k in <dict>".
416422
expr_reg = builder.accept(expr)
417423
target_type = builder.get_dict_key_type(expr)
418-
419-
for_dict = ForDictionaryKeys(builder, index, body_block, loop_exit, line, nested)
424+
for_loop_cls = ForTrueDictionaryKeys if is_true_dict_rprimitive(rtyp) else ForDictionaryKeys
425+
for_dict = for_loop_cls(builder, index, body_block, loop_exit, line, nested)
420426
for_dict.init(expr_reg, target_type)
421427
return for_dict
422428

@@ -498,13 +504,22 @@ def make_for_loop_generator(
498504
for_dict_type: type[ForGenerator] | None = None
499505
if expr.callee.name == "keys":
500506
target_type = builder.get_dict_key_type(expr.callee.expr)
501-
for_dict_type = ForDictionaryKeys
507+
if is_true_dict_rprimitive(rtype):
508+
for_dict_type = ForTrueDictionaryKeys
509+
else:
510+
for_dict_type = ForDictionaryKeys
502511
elif expr.callee.name == "values":
503512
target_type = builder.get_dict_value_type(expr.callee.expr)
504-
for_dict_type = ForDictionaryValues
513+
if is_true_dict_rprimitive(rtype):
514+
for_dict_type = ForTrueDictionaryValues
515+
else:
516+
for_dict_type = ForDictionaryValues
505517
else:
506518
target_type = builder.get_dict_item_type(expr.callee.expr)
507-
for_dict_type = ForDictionaryItems
519+
if is_true_dict_rprimitive(rtype):
520+
for_dict_type = ForTrueDictionaryItems
521+
else:
522+
for_dict_type = ForDictionaryItems
508523
for_dict_gen = for_dict_type(builder, index, body_block, loop_exit, line, nested)
509524
for_dict_gen.init(expr_reg, target_type)
510525
return for_dict_gen
@@ -867,6 +882,7 @@ class ForDictionaryCommon(ForGenerator):
867882

868883
dict_next_op: ClassVar[CFunctionDescription]
869884
dict_iter_op: ClassVar[CFunctionDescription]
885+
dict_size_op: ClassVar[CFunctionDescription] = dict_check_size_op
870886

871887
def need_cleanup(self) -> bool:
872888
# Technically, a dict subclass can raise an unrelated exception
@@ -913,7 +929,7 @@ def gen_step(self) -> None:
913929
line = self.line
914930
# Technically, we don't need a new primitive for this, but it is simpler.
915931
builder.call_c(
916-
dict_check_size_op,
932+
self.dict_size_op,
917933
[builder.read(self.expr_target, line), builder.read(self.size, line)],
918934
line,
919935
)
@@ -991,6 +1007,30 @@ def begin_body(self) -> None:
9911007
builder.assign(target, rvalue, line)
9921008

9931009

1010+
class ForTrueDictionaryKeys(ForDictionaryKeys):
1011+
"""Generate optimized IR for a for loop over dictionary items without type checks."""
1012+
1013+
dict_next_op = true_dict_next_key_op
1014+
dict_iter_op = true_dict_iter_fast_path_op
1015+
dict_size_op = true_dict_check_size_op
1016+
1017+
1018+
class ForTrueDictionaryValues(ForDictionaryValues):
1019+
"""Generate optimized IR for a for loop over dictionary items without type checks."""
1020+
1021+
dict_next_op = true_dict_next_value_op
1022+
dict_iter_op = true_dict_iter_fast_path_op
1023+
dict_size_op = true_dict_check_size_op
1024+
1025+
1026+
class ForTrueDictionaryItems(ForDictionaryItems):
1027+
"""Generate optimized IR for a for loop over dictionary items without type checks."""
1028+
1029+
dict_next_op = true_dict_next_item_op
1030+
dict_iter_op = true_dict_iter_fast_path_op
1031+
dict_size_op = true_dict_check_size_op
1032+
1033+
9941034
class ForRange(ForGenerator):
9951035
"""Generate optimized IR for a for loop over an integer range."""
9961036

mypyc/irbuild/function.py

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,9 @@
5656
from mypyc.ir.rtypes import (
5757
RInstance,
5858
bool_rprimitive,
59-
dict_rprimitive,
6059
int_rprimitive,
6160
object_rprimitive,
61+
true_dict_rprimitive,
6262
)
6363
from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults
6464
from mypyc.irbuild.callable_class import (
@@ -76,7 +76,7 @@
7676
)
7777
from mypyc.irbuild.generator import gen_generator_func, gen_generator_func_body
7878
from mypyc.irbuild.targets import AssignmentTarget
79-
from mypyc.primitives.dict_ops import dict_get_method_with_none, dict_new_op, dict_set_item_op
79+
from mypyc.primitives.dict_ops import dict_new_op, true_dict_get_method_with_none, true_dict_set_item_op
8080
from mypyc.primitives.generic_ops import py_setattr_op
8181
from mypyc.primitives.misc_ops import register_function
8282
from mypyc.primitives.registry import builtin_names
@@ -124,7 +124,7 @@ def transform_decorator(builder: IRBuilder, dec: Decorator) -> None:
124124
if decorated_func is not None:
125125
# Set the callable object representing the decorated function as a global.
126126
builder.primitive_op(
127-
dict_set_item_op,
127+
true_dict_set_item_op,
128128
[builder.load_globals_dict(), builder.load_str(dec.func.name), decorated_func],
129129
decorated_func.line,
130130
)
@@ -797,10 +797,10 @@ def generate_singledispatch_dispatch_function(
797797

798798
arg_type = builder.builder.get_type_of_obj(arg_info.args[0], line)
799799
dispatch_cache = builder.builder.get_attr(
800-
dispatch_func_obj, "dispatch_cache", dict_rprimitive, line
800+
dispatch_func_obj, "dispatch_cache", true_dict_rprimitive, line
801801
)
802802
call_find_impl, use_cache, call_func = BasicBlock(), BasicBlock(), BasicBlock()
803-
get_result = builder.primitive_op(dict_get_method_with_none, [dispatch_cache, arg_type], line)
803+
get_result = builder.primitive_op(true_dict_get_method_with_none, [dispatch_cache, arg_type], line)
804804
is_not_none = builder.translate_is_op(get_result, builder.none_object(), "is not", line)
805805
impl_to_use = Register(object_rprimitive)
806806
builder.add_bool_branch(is_not_none, use_cache, call_find_impl)
@@ -813,7 +813,7 @@ def generate_singledispatch_dispatch_function(
813813
find_impl = builder.load_module_attr_by_fullname("functools._find_impl", line)
814814
registry = load_singledispatch_registry(builder, dispatch_func_obj, line)
815815
uncached_impl = builder.py_call(find_impl, [arg_type, registry], line)
816-
builder.primitive_op(dict_set_item_op, [dispatch_cache, arg_type, uncached_impl], line)
816+
builder.primitive_op(true_dict_set_item_op, [dispatch_cache, arg_type, uncached_impl], line)
817817
builder.assign(impl_to_use, uncached_impl, line)
818818
builder.goto(call_func)
819819

@@ -877,8 +877,8 @@ def gen_dispatch_func_ir(
877877
"""
878878
builder.enter(FuncInfo(fitem, dispatch_name))
879879
setup_callable_class(builder)
880-
builder.fn_info.callable_class.ir.attributes["registry"] = dict_rprimitive
881-
builder.fn_info.callable_class.ir.attributes["dispatch_cache"] = dict_rprimitive
880+
builder.fn_info.callable_class.ir.attributes["registry"] = true_dict_rprimitive
881+
builder.fn_info.callable_class.ir.attributes["dispatch_cache"] = true_dict_rprimitive
882882
builder.fn_info.callable_class.ir.has_dict = True
883883
builder.fn_info.callable_class.ir.needs_getseters = True
884884
generate_singledispatch_callable_class_ctor(builder)
@@ -941,7 +941,7 @@ def add_register_method_to_callable_class(builder: IRBuilder, fn_info: FuncInfo)
941941

942942

943943
def load_singledispatch_registry(builder: IRBuilder, dispatch_func_obj: Value, line: int) -> Value:
944-
return builder.builder.get_attr(dispatch_func_obj, "registry", dict_rprimitive, line)
944+
return builder.builder.get_attr(dispatch_func_obj, "registry", true_dict_rprimitive, line)
945945

946946

947947
def singledispatch_main_func_name(orig_name: str) -> str:
@@ -990,9 +990,9 @@ def maybe_insert_into_registry_dict(builder: IRBuilder, fitem: FuncDef) -> None:
990990
registry = load_singledispatch_registry(builder, dispatch_func_obj, line)
991991
for typ in types:
992992
loaded_type = load_type(builder, typ, None, line)
993-
builder.primitive_op(dict_set_item_op, [registry, loaded_type, to_insert], line)
993+
builder.primitive_op(true_dict_set_item_op, [registry, loaded_type, to_insert], line)
994994
dispatch_cache = builder.builder.get_attr(
995-
dispatch_func_obj, "dispatch_cache", dict_rprimitive, line
995+
dispatch_func_obj, "dispatch_cache", true_dict_rprimitive, line
996996
)
997997
builder.gen_method_call(dispatch_cache, "clear", [], None, line)
998998

0 commit comments

Comments
 (0)