feat: dict literals wip

Author: BobTheBuidler

mypyc/codegen/literals.py

@@ -3,15 +3,15 @@
 from typing import Final, Union
 from typing_extensions import TypeGuard
 
-# Supported Python literal types. All tuple / frozenset items must have supported
+# Supported Python literal types. All tuple / frozenset / dict items must have supported
 # literal types as well, but we can't represent the type precisely.
 LiteralValue = Union[
-    str, bytes, int, bool, float, complex, tuple[object, ...], frozenset[object], None
+    str, bytes, int, bool, float, complex, tuple[object, ...], frozenset[object], dict[object, object], None
 ]
 
 
 def _is_literal_value(obj: object) -> TypeGuard[LiteralValue]:
-    return isinstance(obj, (str, bytes, int, float, complex, tuple, frozenset, type(None)))
+    return isinstance(obj, (str, bytes, int, float, complex, tuple, frozenset, dict, type(None)))
 
 
 # Some literals are singletons and handled specially (None, False and True)
@@ -30,6 +30,7 @@ def __init__(self) -> None:
         self.complex_literals: dict[complex, int] = {}
         self.tuple_literals: dict[tuple[object, ...], int] = {}
         self.frozenset_literals: dict[frozenset[object], int] = {}
+        self.dict_literals: dict[tuple[tuple[object, object], ...], int] = {}
 
     def record_literal(self, value: LiteralValue) -> None:
         """Ensure that the literal value is available in generated code."""
@@ -70,6 +71,17 @@ def record_literal(self, value: LiteralValue) -> None:
                     assert _is_literal_value(item)
                     self.record_literal(item)
                 frozenset_literals[value] = len(frozenset_literals)
+        elif isinstance(value, dict):
+            # Represent dicts as a tuple of sorted (key, value) pairs for uniqueness
+            items = tuple(sorted(value.items()))
+            dict_literals = self.dict_literals
+            if items not in dict_literals:
+                for k, v in items:
+                    assert _is_literal_value(k)
+                    assert _is_literal_value(v)
+                    self.record_literal(k)
+                    self.record_literal(v)
+                dict_literals[items] = len(dict_literals)
         else:
             assert False, "invalid literal: %r" % value
 
@@ -104,6 +116,10 @@ def literal_index(self, value: LiteralValue) -> int:
         n += len(self.tuple_literals)
         if isinstance(value, frozenset):
             return n + self.frozenset_literals[value]
+        n += len(self.frozenset_literals)
+        if isinstance(value, dict):
+            items = tuple(sorted(value.items()))
+            return n + self.dict_literals[items]
         assert False, "invalid literal: %r" % value
 
     def num_literals(self) -> int:
@@ -117,6 +133,7 @@ def num_literals(self) -> int:
             + len(self.complex_literals)
             + len(self.tuple_literals)
             + len(self.frozenset_literals)
+            + len(self.dict_literals)
         )
 
     # The following methods return the C encodings of literal values
@@ -143,6 +160,36 @@ def encoded_tuple_values(self) -> list[str]:
     def encoded_frozenset_values(self) -> list[str]:
         return self._encode_collection_values(self.frozenset_literals)
 
+    def encoded_dict_values(self) -> list[str]:
+        """Encode dict values into a C array.
+
+        The format of the result is like this:
+
+           <number of dicts>
+           <length of the first dict>
+           <literal index of first key>
+           <literal index of first value>
+           ...
+           <literal index of last key>
+           <literal index of last value>
+           <length of the second dict>
+           ...
+        """
+        values = self.dict_literals
+        value_by_index = {index: value for value, index in values.items()}
+        result = []
+        count = len(values)
+        result.append(str(count))
+        for i in range(count):
+            items = value_by_index[i]
+            result.append(str(len(items)))
+            for k, v in items:
+                index_k = self.literal_index(k)
+                index_v = self.literal_index(v)
+                result.append(str(index_k))
+                result.append(str(index_v))
+        return result
+
     def _encode_collection_values(
         self, values: dict[tuple[object, ...], int] | dict[frozenset[object], int]
     ) -> list[str]:

mypyc/irbuild/expression.py

@@ -59,9 +59,11 @@
     BasicBlock,
     Call,
     ComparisonOp,
+    InitStatic,
     Integer,
     LoadAddress,
     LoadLiteral,
+    LoadStatic,
     PrimitiveDescription,
     RaiseStandardError,
     Register,
@@ -73,6 +75,7 @@
     RInstance,
     RTuple,
     bool_rprimitive,
+    dict_rprimitive,
     int_rprimitive,
     is_fixed_width_rtype,
     is_int_rprimitive,
@@ -100,7 +103,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_get_item_op, dict_new_op, dict_set_item_op, dict_template_copy_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
@@ -1009,8 +1012,53 @@ def _visit_tuple_display(builder: IRBuilder, expr: TupleExpr) -> Value:
     return builder.primitive_op(list_tuple_op, [val_as_list], expr.line)
 
 
+def dict_literal_values(builder: IRBuilder, items: Sequence[tuple[Expression | None, Expression]], line: int) -> "Value | None":
+    """Try to extract a constant dict from a dict literal.
+
+    If all keys and values are deeply immutable and constant, build a static dict at module init,
+    register it as a static, and return a LoadStatic for it. Otherwise, return None.
+    """
+    result = {}
+    for key_expr, value_expr in items:
+        if key_expr is None:
+            # ** unpacking, not a literal
+            # TODO: if ** is unpacking a dict literal we can use that, we just need logic
+            return None
+        key = constant_fold_expr(builder, key_expr)
+        if key is None:
+            return None
+        value = constant_fold_expr(builder, value_expr)
+        if value is None:
+            return None
+        result[key] = value
+
+    dict_reg = builder.call_c(dict_new_op, [], line)
+    for k, v in result.items():
+        key = builder.load_literal_value(k)
+        value = builder.load_literal_value(v)
+        builder.primitive_op(dict_set_item_op, [dict_reg, key, value], line)
+
+    # Register as a static with a unique name
+    static_name = f"__mypyc_dict_template__{abs(hash(tuple(sorted(result.items()))))}"
+
+    if _static_dicts.get(static_name) is None:
+        _static_dicts[static_name] = builder.add(InitStatic(dict_reg, static_name, builder.module_name))
+    
+    return builder.add(LoadStatic(dict_rprimitive, static_name, builder.module_name))
+
 def transform_dict_expr(builder: IRBuilder, expr: DictExpr) -> Value:
-    """First accepts all keys and values, then makes a dict out of them."""
+    """First accepts all keys and values, then makes a dict out of them.
+
+    Optimization: If all keys and values are deeply immutable, emit a static template dict
+    and at runtime use PyDict_Copy to return a fresh dict.
+    """
+    # Try to constant fold the dict and get a static Value
+    template = dict_literal_values(builder, expr.items, expr.line)
+    if template is not None:
+        # At runtime, return PyDict_Copy(template)
+        return builder.call_c(dict_template_copy_op, [template], expr.line)
+
+    # Fallback: build dict at runtime as before
     key_value_pairs = []
     for key_expr, value_expr in expr.items:
         key = builder.accept(key_expr) if key_expr is not None else None
@@ -1020,6 +1068,8 @@ def transform_dict_expr(builder: IRBuilder, expr: DictExpr) -> Value:
     return builder.builder.make_dict(key_value_pairs, expr.line)
 
 
+_static_dicts = {}
+
 def transform_set_expr(builder: IRBuilder, expr: SetExpr) -> Value:
     return _visit_display(
         builder, expr.items, builder.new_set_op, set_add_op, set_update_op, expr.line, False

mypyc/primitives/dict_ops.py

@@ -62,6 +62,13 @@
     priority=2,
 )
 
+dict_template_copy_op = custom_op(
+    arg_types=[dict_rprimitive],
+    return_type=dict_rprimitive,
+    c_function_name="PyDict_Copy",
+    error_kind=ERR_MAGIC,
+)
+
 # Generic one-argument dict constructor: dict(obj)
 dict_copy = function_op(
     name="builtins.dict",

mypyc/test-data/run-dicts.test

@@ -368,3 +368,107 @@ class subc(dict[Any, Any]):
 [file userdefineddict.py]
 class dict:
     pass
+
+[case testDictConstantFolding]
+# Test constant folding of dict literals with deeply immutable keys/values
+def get_headers():
+    return {"k0": "v0", "k1": "v1"}
+
+def get_numbers():
+    return {1: 2, 3: 4}
+
+def get_tuple_key():
+    return {(1, 2): "a", (3, 4): "b"}
+
+def get_frozenset_key():
+    return {frozenset({1, 2}): "x"}
+
+def get_nested_tuple():
+    return {(1, (2, 3)): "ok"}
+
+def get_bool_none():
+    return {True: None, False: 1}
+
+def get_bytes():
+    return {b"x": b"y"}
+
+def get_mutable_value():
+    return {"a": []}
+
+def get_mutable_key():
+    return {(1, []): "bad"}
+
+def get_nested_mutable():
+    return {(1, (2, [])): "bad"}
+
+def get_unhashable_key():
+    return {[1, 2]: "bad"}
+
+def get_tuple_of_frozenset():
+    return {(1, frozenset({2, 3})): "ok"}
+
+def get_tuple_of_tuple():
+    return {(1, (2, (3, 4))): "ok"}
+
+def get_frozenset_of_tuple():
+    return {frozenset({(1, 2), (3, 4)}): "ok"}
+
+class C: pass
+
+def get_user_defined():
+    return {C(): 1}
+
+def test_mutation_independence():
+    d1 = get_headers()
+    d2 = get_headers()
+    d1["k0"] = "changed"
+    assert d2["k0"] == "v0"
+    d2["k1"] = "changed2"
+    assert d1["k1"] == "v1"
+
+def test_ineligible_not_folded():
+    # These should not be constant folded
+    assert get_mutable_value()["a"] == []
+    try:
+        get_mutable_key()
+    except TypeError:
+        pass
+    else:
+        assert False
+    try:
+        get_nested_mutable()
+    except TypeError:
+        pass
+    else:
+        assert False
+    try:
+        get_unhashable_key()
+    except TypeError:
+        pass
+    else:
+        assert False
+    # User-defined class instance as key should not be folded
+    d = get_user_defined()
+    assert list(d.values()) == [1]
+
+def test_eligible_folded():
+    assert get_headers() == {"k0": "v0", "k1": "v1"}
+    assert get_numbers() == {1: 2, 3: 4}
+    assert get_tuple_key() == {(1, 2): "a", (3, 4): "b"}
+    assert get_frozenset_key() == {frozenset({1, 2}): "x"}
+    assert get_nested_tuple() == {(1, (2, 3)): "ok"}
+    assert get_bool_none() == {True: None, False: 1}
+    assert get_bytes() == {b"x": b"y"}
+    assert get_tuple_of_frozenset() == {(1, frozenset({2, 3})): "ok"}
+    assert get_tuple_of_tuple() == {(1, (2, (3, 4))): "ok"}
+    assert get_frozenset_of_tuple() == {frozenset({(1, 2), (3, 4)}): "ok"}
+
+[file driver.py]
+from native import (
+    test_mutation_independence,
+    test_ineligible_not_folded,
+    test_eligible_folded,
+)
+test_mutation_independence()
+test_ineligible_not_folded()
+test_eligible_folded()