[mypyc] feat: constant folding for DictExpr [1/1]

mypyc/analysis/ircheck.py

@@ -285,6 +285,16 @@ def check_frozenset_items_valid_literals(self, op: LoadLiteral, s: frozenset[obj
             else:
                 self.fail(op, f"Invalid type for item of frozenset literal: {type(x)})")
 
+    def check_dict_items_valid_literals(self, op: LoadLiteral, d: dict[object, object]) -> None:
+        valid_types = (str, bytes, bool, int, float, complex)
+        for k, v in d.items():
+            # Acceptable key types: str, bytes, bool, int, float, complex
+            if not isinstance(k, valid_types):
+                self.fail(op, f"Invalid type for key of dict literal: {type(k)})")
+            # Acceptable value types: str, bytes, bool, int, float, complex
+            if not isinstance(v, valid_types):
+                self.fail(op, f"Invalid type for value of dict literal: {type(v)})")
+
     def visit_load_literal(self, op: LoadLiteral) -> None:
         expected_type = None
         if op.value is None:
@@ -309,6 +319,9 @@ def visit_load_literal(self, op: LoadLiteral) -> None:
             # it's a set (when it's really a frozenset).
             expected_type = "builtins.set"
             self.check_frozenset_items_valid_literals(op, op.value)
+        elif isinstance(op.value, dict):
+            expected_type = "builtins.dict"
+            self.check_dict_items_valid_literals(op, op.value)
 
         assert expected_type is not None, "Missed a case for LoadLiteral check"
 

mypyc/codegen/emitmodule.py

@@ -708,6 +708,9 @@ def generate_literal_tables(self) -> None:
         # Descriptions of frozenset literals
         init_frozenset = c_array_initializer(literals.encoded_frozenset_values())
         self.declare_global("const int []", "CPyLit_FrozenSet", initializer=init_frozenset)
+        # Descriptions of dict literals
+        init_dict = c_array_initializer(literals.encoded_dict_values())
+        self.declare_global("const int []", "CPyLit_Dict", initializer=init_dict)
 
     def generate_export_table(self, decl_emitter: Emitter, code_emitter: Emitter) -> None:
         """Generate the declaration and definition of the group's export struct.
@@ -926,7 +929,7 @@ def generate_globals_init(self, emitter: Emitter) -> None:
         for symbol, fixup in self.simple_inits:
             emitter.emit_line(f"{symbol} = {fixup};")
 
-        values = "CPyLit_Str, CPyLit_Bytes, CPyLit_Int, CPyLit_Float, CPyLit_Complex, CPyLit_Tuple, CPyLit_FrozenSet"
+        values = "CPyLit_Str, CPyLit_Bytes, CPyLit_Int, CPyLit_Float, CPyLit_Complex, CPyLit_Tuple, CPyLit_FrozenSet, CPyLit_Dict"
         emitter.emit_lines(
             f"if (CPyStatics_Initialize(CPyStatics, {values}) < 0) {{", "return -1;", "}"
         )

mypyc/codegen/literals.py

@@ -3,15 +3,24 @@
 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 +39,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 +80,16 @@ 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):
+            items = self.make_dict_literal_key(value)
+            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,8 +124,18 @@ 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):
+            key = self.make_dict_literal_key(value)
+            return n + self.dict_literals[key]
         assert False, "invalid literal: %r" % value
 
+    def make_dict_literal_key(
+        self, value: dict[LiteralValue, LiteralValue]
+    ) -> tuple[tuple[LiteralValue, LiteralValue]]:
+        """Make a unique key for a literal dict."""
+        return tuple(value.items())
+
     def num_literals(self) -> int:
         # The first three are for None, True and False
         return (
@@ -117,6 +147,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 +174,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/ir/ops.py

@@ -812,6 +812,10 @@ class LoadLiteral(RegisterOp):
 
     Tuple / frozenset literals must contain only valid literal values as items.
 
+    Dict literals must contain only literal keys and literal values.
+    Due to their mutability, dict literals will be copied from the main template
+    at each use.
+
     NOTE: You can use this to load boxed (Python) int objects. Use
           Integer to load unboxed, tagged integers or fixed-width,
           low-level integers.

mypyc/irbuild/expression.py

@@ -73,6 +73,7 @@
     RInstance,
     RTuple,
     bool_rprimitive,
+    dict_rprimitive,
     int_rprimitive,
     is_fixed_width_rtype,
     is_int_rprimitive,
@@ -100,7 +101,12 @@
 )
 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_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 +1015,46 @@ 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
+) -> dict | None:
+    """Try to extract a constant dict from a dict literal, recursively staticizing nested dicts.
+
+    If all keys and values are deeply immutable and constant (including nested dicts as values),
+    return the Python dict value. 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
+        # Recursively staticize dict values
+        value = constant_fold_expr(builder, value_expr)
+        if value is None:
+            return None
+        result[key] = value
+
+    return result or None
+
+
 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
+    static_dict = dict_literal_values(builder, expr.items, expr.line)
+    if static_dict is not None:
+        # Register the static dict and return a copy at runtime
+        static_val = builder.add(LoadLiteral(static_dict, dict_rprimitive))
+        return builder.call_c(dict_template_copy_op, [static_val], expr.line)
+
+    # If that fails, build dict at runtime
     key_value_pairs = []
     for key_expr, value_expr in expr.items:
         key = builder.accept(key_expr) if key_expr is not None else None

mypyc/lib-rt/CPy.h

@@ -931,7 +931,8 @@ int CPyStatics_Initialize(PyObject **statics,
                           const double *floats,
                           const double *complex_numbers,
                           const int *tuples,
-                          const int *frozensets);
+                          const int *frozensets,
+                          const int *dicts);
 PyObject *CPy_Super(PyObject *builtins, PyObject *self);
 PyObject *CPy_CallReverseOpMethod(PyObject *left, PyObject *right, const char *op,
                                   _Py_Identifier *method);

mypyc/lib-rt/misc_ops.c

@@ -615,7 +615,8 @@ int CPyStatics_Initialize(PyObject **statics,
                           const double *floats,
                           const double *complex_numbers,
                           const int *tuples,
-                          const int *frozensets) {
+                          const int *frozensets,
+                          const int *dicts) {
     PyObject **result = statics;
     // Start with some hard-coded values
     *result++ = Py_None;
@@ -733,6 +734,31 @@ int CPyStatics_Initialize(PyObject **statics,
             *result++ = obj;
         }
     }
+    if (dicts) {
+        int num = *dicts++;
+        while (num-- > 0) {
+            int num_items = *dicts++;
+            PyObject *obj = PyDict_New();
+            if (obj == NULL) {
+                return -1;
+            }
+            for (int i = 0; i < num_items; i++) {
+                PyObject *key = statics[*dicts++];
+                PyObject *value = statics[*dicts++];
+                Py_INCREF(key);
+                Py_INCREF(value);
+                if (PyDict_SetItem(obj, key, value) == -1) {
+                    Py_DECREF(key);
+                    Py_DECREF(value);
+                    Py_DECREF(obj);
+                    return -1;
+                }
+                Py_DECREF(key);
+                Py_DECREF(value);
+            }
+            *result++ = obj;
+        }
+    }
     return 0;
 }
 

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",