[mypyc] Inline math literals

mypyc/codegen/emitfunc.py

@@ -775,6 +775,8 @@ def reg(self, reg: Value) -> str:
                 return "INFINITY"
             elif r == "-inf":
                 return "-INFINITY"
+            elif r == "nan":
+                return "NAN"
             return r
         else:
             return self.emitter.reg(reg)

mypyc/codegen/literals.py

@@ -265,6 +265,8 @@ def float_to_c(x: float) -> str:
         return "INFINITY"
     elif s == "-inf":
         return "-INFINITY"
+    elif s == "nan":
+        return "NAN"
     return s
 
 

mypyc/irbuild/builder.py

@@ -301,6 +301,9 @@ def load_bytes_from_str_literal(self, value: str) -> Value:
     def load_int(self, value: int) -> Value:
         return self.builder.load_int(value)
 
+    def load_float(self, value: float) -> Value:
+        return self.builder.load_float(value)
+
     def unary_op(self, lreg: Value, expr_op: str, line: int) -> Value:
         return self.builder.unary_op(lreg, expr_op, line)
 

mypyc/irbuild/expression.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import math
 from typing import Callable, Sequence
 
 from mypy.nodes import (
@@ -128,6 +129,10 @@ def transform_name_expr(builder: IRBuilder, expr: NameExpr) -> Value:
     if fullname == "builtins.False":
         return builder.false()
 
+    math_literal = transform_math_literal(builder, fullname)
+    if math_literal is not None:
+        return math_literal
+
     if isinstance(expr.node, Var) and expr.node.is_final:
         value = builder.emit_load_final(
             expr.node,
@@ -192,6 +197,10 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
         if value is not None:
             return value
 
+    math_literal = transform_math_literal(builder, expr.fullname)
+    if math_literal is not None:
+        return math_literal
+
     if isinstance(expr.node, MypyFile) and expr.node.fullname in builder.imports:
         return builder.load_module(expr.node.fullname)
 
@@ -1043,3 +1052,18 @@ def transform_assignment_expr(builder: IRBuilder, o: AssignmentExpr) -> Value:
     target = builder.get_assignment_target(o.target)
     builder.assign(target, value, o.line)
     return value
+
+
+def transform_math_literal(builder: IRBuilder, fullname: str) -> Value | None:
+    if fullname == "math.e":
+        return builder.load_float(math.e)
+    if fullname == "math.pi":
+        return builder.load_float(math.pi)
+    if fullname == "math.inf":
+        return builder.load_float(math.inf)
+    if fullname == "math.nan":
+        return builder.load_float(math.nan)
+    if fullname == "math.tau":
+        return builder.load_float(math.tau)
+
+    return None

mypyc/test-data/run-math.test

@@ -4,6 +4,7 @@
 from typing import Any, Callable
 from typing_extensions import Final
 import math
+from math import pi, e, tau, inf, nan
 from testutil import assertRaises, float_vals, assertDomainError, assertMathRangeError
 
 pymath: Any = math
@@ -86,3 +87,20 @@ def test_isinf() -> None:
 def test_isnan() -> None:
     for x in float_vals:
         assert repr(math.isnan(x)) == repr(pymath.isnan(x))
+
+
+def test_pi_is_inlined_correctly() -> None:
+    assert math.pi == pi == 3.141592653589793
+
+def test_e_is_inlined_correctly() -> None:
+    assert math.e == e == 2.718281828459045
+
+def test_tau_is_inlined_correctly() -> None:
+    assert math.tau == tau == 6.283185307179586
+
+def test_inf_is_inlined_correctly() -> None:
+    assert math.inf == inf == float("inf")
+
+def test_nan_is_inlined_correctly() -> None:
+    assert math.isnan(math.nan)
+    assert math.isnan(nan)