Constant fold initializers of final variables (#14283)

Now mypy can figure out the values of final variables even if the
initializer has some operations on constant values:
```
A: Final = 2  # This has always worked
A: Final = -(1 << 2)  # This is now supported
B: Final = 'x' + 'y'  # This also now works
```
Currently we support integer arithmetic and bitwise operations, and
string concatenation.

This can be useful with literal types, but my main goal was to improve
constant folding in mypyc. In particular, this helps constant folding
with native ints in cases like these:
```
FLAG1: Final = 1 << 4
FLAG2: Final = 1 << 5

def f() -> i64:
    return FLAG1 | FLAG2  # Can now be constant folded
```

We still have another constant folding pass in mypyc, since it does some
things more aggressively (e.g. it constant folds some member expression
references).

Work on mypyc/mypyc#772.

Also helps with mypyc/mypyc#862.

Author: JukkaL

mypy/constant_fold.py

@@ -0,0 +1,116 @@
+"""Constant folding of expressions.
+
+For example, 3 + 5 can be constant folded into 8.
+"""
+
+from __future__ import annotations
+
+from typing import Union
+from typing_extensions import Final
+
+from mypy.nodes import Expression, FloatExpr, IntExpr, NameExpr, OpExpr, StrExpr, UnaryExpr, Var
+
+# All possible result types of constant folding
+ConstantValue = Union[int, bool, float, str]
+CONST_TYPES: Final = (int, bool, float, str)
+
+
+def constant_fold_expr(expr: Expression, cur_mod_id: str) -> ConstantValue | None:
+    """Return the constant value of an expression for supported operations.
+
+    Among other things, support int arithmetic and string
+    concatenation. For example, the expression 3 + 5 has the constant
+    value 8.
+
+    Also bind simple references to final constants defined in the
+    current module (cur_mod_id). Binding to references is best effort
+    -- we don't bind references to other modules. Mypyc trusts these
+    to be correct in compiled modules, so that it can replace a
+    constant expression (or a reference to one) with the statically
+    computed value. We don't want to infer constant values based on
+    stubs, in particular, as these might not match the implementation
+    (due to version skew, for example).
+
+    Return None if unsuccessful.
+    """
+    if isinstance(expr, IntExpr):
+        return expr.value
+    if isinstance(expr, StrExpr):
+        return expr.value
+    if isinstance(expr, FloatExpr):
+        return expr.value
+    elif isinstance(expr, NameExpr):
+        if expr.name == "True":
+            return True
+        elif expr.name == "False":
+            return False
+        node = expr.node
+        if (
+            isinstance(node, Var)
+            and node.is_final
+            and node.fullname.rsplit(".", 1)[0] == cur_mod_id
+        ):
+            value = node.final_value
+            if isinstance(value, (CONST_TYPES)):
+                return value
+    elif isinstance(expr, OpExpr):
+        left = constant_fold_expr(expr.left, cur_mod_id)
+        right = constant_fold_expr(expr.right, cur_mod_id)
+        if isinstance(left, int) and isinstance(right, int):
+            return constant_fold_binary_int_op(expr.op, left, right)
+        elif isinstance(left, str) and isinstance(right, str):
+            return constant_fold_binary_str_op(expr.op, left, right)
+    elif isinstance(expr, UnaryExpr):
+        value = constant_fold_expr(expr.expr, cur_mod_id)
+        if isinstance(value, int):
+            return constant_fold_unary_int_op(expr.op, value)
+    return None
+
+
+def constant_fold_binary_int_op(op: str, left: int, right: int) -> int | None:
+    if op == "+":
+        return left + right
+    if op == "-":
+        return left - right
+    elif op == "*":
+        return left * right
+    elif op == "//":
+        if right != 0:
+            return left // right
+    elif op == "%":
+        if right != 0:
+            return left % right
+    elif op == "&":
+        return left & right
+    elif op == "|":
+        return left | right
+    elif op == "^":
+        return left ^ right
+    elif op == "<<":
+        if right >= 0:
+            return left << right
+    elif op == ">>":
+        if right >= 0:
+            return left >> right
+    elif op == "**":
+        if right >= 0:
+            ret = left**right
+            assert isinstance(ret, int)
+            return ret
+    return None
+
+
+def constant_fold_unary_int_op(op: str, value: int) -> int | None:
+    if op == "-":
+        return -value
+    elif op == "~":
+        return ~value
+    elif op == "+":
+        return value
+    return None
+
+
+def constant_fold_binary_str_op(op: str, left: str, right: str) -> str | None:
+    if op == "+":
+        return left + right
+    return None

mypy/semanal.py

@@ -55,6 +55,7 @@
 from typing_extensions import Final, TypeAlias as _TypeAlias
 
 from mypy import errorcodes as codes, message_registry
+from mypy.constant_fold import constant_fold_expr
 from mypy.errorcodes import ErrorCode
 from mypy.errors import Errors, report_internal_error
 from mypy.exprtotype import TypeTranslationError, expr_to_unanalyzed_type
@@ -91,7 +92,6 @@
     AwaitExpr,
     Block,
     BreakStmt,
-    BytesExpr,
     CallExpr,
     CastExpr,
     ClassDef,
@@ -108,7 +108,6 @@
     Expression,
     ExpressionStmt,
     FakeExpression,
-    FloatExpr,
     ForStmt,
     FuncBase,
     FuncDef,
@@ -121,7 +120,6 @@
     ImportBase,
     ImportFrom,
     IndexExpr,
-    IntExpr,
     LambdaExpr,
     ListComprehension,
     ListExpr,
@@ -250,7 +248,6 @@
     FunctionLike,
     Instance,
     LiteralType,
-    LiteralValue,
     NoneType,
     Overloaded,
     Parameters,
@@ -3138,7 +3135,8 @@ def store_final_status(self, s: AssignmentStmt) -> None:
                 node = s.lvalues[0].node
                 if isinstance(node, Var):
                     node.is_final = True
-                    node.final_value = self.unbox_literal(s.rvalue)
+                    if s.type:
+                        node.final_value = constant_fold_expr(s.rvalue, self.cur_mod_id)
                     if self.is_class_scope() and (
                         isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs
                     ):
@@ -3198,13 +3196,6 @@ def flatten_lvalues(self, lvalues: list[Expression]) -> list[Expression]:
                 res.append(lv)
         return res
 
-    def unbox_literal(self, e: Expression) -> int | float | bool | str | None:
-        if isinstance(e, (IntExpr, FloatExpr, StrExpr)):
-            return e.value
-        elif isinstance(e, NameExpr) and e.name in ("True", "False"):
-            return True if e.name == "True" else False
-        return None
-
     def process_type_annotation(self, s: AssignmentStmt) -> None:
         """Analyze type annotation or infer simple literal type."""
         if s.type:
@@ -3259,39 +3250,33 @@ def is_annotated_protocol_member(self, s: AssignmentStmt) -> bool:
 
     def analyze_simple_literal_type(self, rvalue: Expression, is_final: bool) -> Type | None:
         """Return builtins.int if rvalue is an int literal, etc.
-        If this is a 'Final' context, we return "Literal[...]" instead."""
-        if self.options.semantic_analysis_only or self.function_stack:
-            # Skip this if we're only doing the semantic analysis pass.
-            # This is mostly to avoid breaking unit tests.
-            # Also skip inside a function; this is to avoid confusing
+
+        If this is a 'Final' context, we return "Literal[...]" instead.
+        """
+        if self.function_stack:
+            # Skip inside a function; this is to avoid confusing
             # the code that handles dead code due to isinstance()
             # inside type variables with value restrictions (like
             # AnyStr).
             return None
-        if isinstance(rvalue, FloatExpr):
-            return self.named_type_or_none("builtins.float")
-
-        value: LiteralValue | None = None
-        type_name: str | None = None
-        if isinstance(rvalue, IntExpr):
-            value, type_name = rvalue.value, "builtins.int"
-        if isinstance(rvalue, StrExpr):
-            value, type_name = rvalue.value, "builtins.str"
-        if isinstance(rvalue, BytesExpr):
-            value, type_name = rvalue.value, "builtins.bytes"
-
-        if type_name is not None:
-            assert value is not None
-            typ = self.named_type_or_none(type_name)
-            if typ and is_final:
-                return typ.copy_modified(
-                    last_known_value=LiteralType(
-                        value=value, fallback=typ, line=typ.line, column=typ.column
-                    )
-                )
-            return typ
 
-        return None
+        value = constant_fold_expr(rvalue, self.cur_mod_id)
+        if value is None:
+            return None
+
+        if isinstance(value, bool):
+            type_name = "builtins.bool"
+        elif isinstance(value, int):
+            type_name = "builtins.int"
+        elif isinstance(value, str):
+            type_name = "builtins.str"
+        elif isinstance(value, float):
+            type_name = "builtins.float"
+
+        typ = self.named_type_or_none(type_name)
+        if typ and is_final:
+            return typ.copy_modified(last_known_value=LiteralType(value=value, fallback=typ))
+        return typ
 
     def analyze_alias(
         self, name: str, rvalue: Expression, allow_placeholder: bool = False
@@ -3827,6 +3812,14 @@ def store_declared_types(self, lvalue: Lvalue, typ: Type) -> None:
                 var = lvalue.node
                 var.type = typ
                 var.is_ready = True
+                typ = get_proper_type(typ)
+                if (
+                    var.is_final
+                    and isinstance(typ, Instance)
+                    and typ.last_known_value
+                    and (not self.type or not self.type.is_enum)
+                ):
+                    var.final_value = typ.last_known_value.value
             # If node is not a variable, we'll catch it elsewhere.
         elif isinstance(lvalue, TupleExpr):
             typ = get_proper_type(typ)

mypy/types.py

@@ -67,7 +67,10 @@
 # Note: Although "Literal[None]" is a valid type, we internally always convert
 # such a type directly into "None". So, "None" is not a valid parameter of
 # LiteralType and is omitted from this list.
-LiteralValue: _TypeAlias = Union[int, str, bool]
+#
+# Note: Float values are only used internally. They are not accepted within
+# Literal[...].
+LiteralValue: _TypeAlias = Union[int, str, bool, float]
 
 
 # If we only import type_visitor in the middle of the file, mypy

mypyc/irbuild/constant_fold.py

@@ -1,13 +1,23 @@
 """Constant folding of IR values.
 
 For example, 3 + 5 can be constant folded into 8.
+
+This is mostly like mypy.constant_fold, but we can bind some additional
+NameExpr and MemberExpr references here, since we have more knowledge
+about which definitions can be trusted -- we constant fold only references
+to other compiled modules in the same compilation unit.
 """
 
 from __future__ import annotations
 
 from typing import Union
 from typing_extensions import Final
 
+from mypy.constant_fold import (
+    constant_fold_binary_int_op,
+    constant_fold_binary_str_op,
+    constant_fold_unary_int_op,
+)
 from mypy.nodes import Expression, IntExpr, MemberExpr, NameExpr, OpExpr, StrExpr, UnaryExpr, Var
 from mypyc.irbuild.builder import IRBuilder
 
@@ -51,52 +61,3 @@ def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue |
         if isinstance(value, int):
             return constant_fold_unary_int_op(expr.op, value)
     return None
-
-
-def constant_fold_binary_int_op(op: str, left: int, right: int) -> int | None:
-    if op == "+":
-        return left + right
-    if op == "-":
-        return left - right
-    elif op == "*":
-        return left * right
-    elif op == "//":
-        if right != 0:
-            return left // right
-    elif op == "%":
-        if right != 0:
-            return left % right
-    elif op == "&":
-        return left & right
-    elif op == "|":
-        return left | right
-    elif op == "^":
-        return left ^ right
-    elif op == "<<":
-        if right >= 0:
-            return left << right
-    elif op == ">>":
-        if right >= 0:
-            return left >> right
-    elif op == "**":
-        if right >= 0:
-            ret = left**right
-            assert isinstance(ret, int)
-            return ret
-    return None
-
-
-def constant_fold_unary_int_op(op: str, value: int) -> int | None:
-    if op == "-":
-        return -value
-    elif op == "~":
-        return ~value
-    elif op == "+":
-        return value
-    return None
-
-
-def constant_fold_binary_str_op(op: str, left: str, right: str) -> str | None:
-    if op == "+":
-        return left + right
-    return None

mypyc/test-data/irbuild-basic.test

@@ -3273,7 +3273,7 @@ L2:
 [case testFinalStaticInt]
 from typing import Final
 
-x: Final = 1 + 1
+x: Final = 1 + int()
 
 def f() -> int:
     return x - 1

mypyc/test-data/irbuild-constant-fold.test

@@ -205,23 +205,13 @@ Y: Final = 2 + 4
 
 def f() -> None:
     a = X + 1
-    # TODO: Constant fold this as well
     a = Y + 1
 [out]
 def f():
-    a, r0 :: int
-    r1 :: bool
-    r2 :: int
+    a :: int
 L0:
     a = 12
-    r0 = __main__.Y :: static
-    if is_error(r0) goto L1 else goto L2
-L1:
-    r1 = raise NameError('value for final name "Y" was not set')
-    unreachable
-L2:
-    r2 = CPyTagged_Add(r0, 2)
-    a = r2
+    a = 14
     return 1
 
 [case testIntConstantFoldingClassFinal]