[mypyc] Inline math literals (#15324)

This PR inlines math literals such as `math.pi` and `math.e`. Previously
Using `math.pi` would emit a getattr, unbox, error check, then box op
code (if we need a boxed value), but now it just generates a literal directly
(that is sometimes boxed).

Closes mypyc/mypyc#985

Author: dosisod

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/irbuild-math.test

@@ -0,0 +1,64 @@
+[case testMathLiteralsAreInlined]
+import math
+from math import pi, e, tau, inf, nan
+
+def f1() -> float:
+    return pi
+
+def f2() -> float:
+    return math.pi
+
+def f3() -> float:
+    return math.e
+
+def f4() -> float:
+    return math.e
+
+def f5() -> float:
+    return math.tau
+
+def f6() -> float:
+    return math.tau
+
+def f7() -> float:
+    return math.inf
+def f8() -> float:
+    return math.inf
+
+def f9() -> float:
+    return math.nan
+
+def f10() -> float:
+    return math.nan
+
+[out]
+def f1():
+L0:
+    return 3.141592653589793
+def f2():
+L0:
+    return 3.141592653589793
+def f3():
+L0:
+    return 2.718281828459045
+def f4():
+L0:
+    return 2.718281828459045
+def f5():
+L0:
+    return 6.283185307179586
+def f6():
+L0:
+    return 6.283185307179586
+def f7():
+L0:
+    return inf
+def f8():
+L0:
+    return inf
+def f9():
+L0:
+    return nan
+def f10():
+L0:
+    return nan

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)

mypyc/test/test_irbuild.py

@@ -49,6 +49,7 @@
     "irbuild-singledispatch.test",
     "irbuild-constant-fold.test",
     "irbuild-glue-methods.test",
+    "irbuild-math.test",
 ]
 
 if sys.version_info >= (3, 10):

test-data/unit/lib-stub/math.pyi

@@ -1,4 +1,8 @@
 pi: float
+e: float
+tau: float
+inf: float
+nan: float
 def sqrt(__x: float) -> float: ...
 def sin(__x: float) -> float: ...
 def cos(__x: float) -> float: ...